|
22 | 22 | print("Usage: python {} IMAGE ANNO OUTPUT".format(sys.argv[0])) |
23 | 23 | print("") |
24 | 24 | print("IMAGE and ANNO are inputs and OUTPUT is where the result should be written.") |
| 25 | + print("If there's at least one single full-black pixel in ANNO, black is assumed to mean unknown.") |
25 | 26 | sys.exit(1) |
26 | 27 |
|
27 | 28 | fn_im = sys.argv[1] |
|
41 | 42 | # Note that all-black, i.e. the value 0 for background will stay 0. |
42 | 43 | colors, labels = np.unique(anno_lbl, return_inverse=True) |
43 | 44 |
|
44 | | -# And create a mapping back from the labels to 32bit integer colors. |
45 | 45 | # But remove the all-0 black, that won't exist in the MAP! |
46 | | -colors = colors[1:] |
| 46 | +HAS_UNK = 0 in colors |
| 47 | +if HAS_UNK: |
| 48 | + print("Found a full-black pixel in annotation image, assuming it means 'unknown' label!") |
| 49 | + colors = colors[1:] |
| 50 | +#else: |
| 51 | +# print("No single full-black pixel found in annotation image. Assuming there's no 'unknown' label!") |
| 52 | + |
| 53 | +# And create a mapping back from the labels to 32bit integer colors. |
47 | 54 | colorize = np.empty((len(colors), 3), np.uint8) |
48 | 55 | colorize[:,0] = (colors & 0x0000FF) |
49 | 56 | colorize[:,1] = (colors & 0x00FF00) >> 8 |
|
52 | 59 | # Compute the number of classes in the label image. |
53 | 60 | # We subtract one because the number shouldn't include the value 0 which stands |
54 | 61 | # for "unknown" or "unsure". |
55 | | -n_labels = len(set(labels.flat)) - 1 |
56 | | -print(n_labels, " labels and \"unknown\" 0: ", set(labels.flat)) |
| 62 | +n_labels = len(set(labels.flat)) - int(HAS_UNK) |
| 63 | +print(n_labels, " labels", (" plus \"unknown\" 0: " if HAS_UNK else ""), set(labels.flat)) |
57 | 64 |
|
58 | 65 | ########################### |
59 | 66 | ### Setup the CRF model ### |
|
67 | 74 | d = dcrf.DenseCRF2D(img.shape[1], img.shape[0], n_labels) |
68 | 75 |
|
69 | 76 | # get unary potentials (neg log probability) |
70 | | - U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=True) |
| 77 | + U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=HAS_UNK) |
71 | 78 | d.setUnaryEnergy(U) |
72 | 79 |
|
73 | 80 | # This adds the color-independent term, features are the locations only. |
|
86 | 93 | d = dcrf.DenseCRF(img.shape[1] * img.shape[0], n_labels) |
87 | 94 |
|
88 | 95 | # get unary potentials (neg log probability) |
89 | | - U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=True) |
| 96 | + U = unary_from_labels(labels, n_labels, gt_prob=0.7, zero_unsure=HAS_UNK) |
90 | 97 | d.setUnaryEnergy(U) |
91 | 98 |
|
92 | 99 | # This creates the color-independent features and then add them to the CRF |
|
114 | 121 | MAP = np.argmax(Q, axis=0) |
115 | 122 |
|
116 | 123 | # Convert the MAP (labels) back to the corresponding colors and save the image. |
| 124 | +# Note that there is no "unknown" here anymore, no matter what we had at first. |
117 | 125 | MAP = colorize[MAP,:] |
118 | | -imsave(fn_output, MAP.reshape(img.shape)) |
| 126 | +imwrite(fn_output, MAP.reshape(img.shape)) |
119 | 127 |
|
120 | 128 | # Just randomly manually run inference iterations |
121 | 129 | Q, tmp1, tmp2 = d.startInference() |
|
0 commit comments