Args: binary_mask (ndarray): numpy array of shape (H, W), where H is the image height and W is the image width. Each value in the array is either a 0 or 1 value of uint8 type. color: color of the mask. Refer to `matplotlib.colors` for
(
self, binary_mask, color=None, *, edge_color=None, text=None, alpha=0.5, area_threshold=0
)
| 1003 | return self.output |
| 1004 | |
| 1005 | def draw_binary_mask( |
| 1006 | self, binary_mask, color=None, *, edge_color=None, text=None, alpha=0.5, area_threshold=0 |
| 1007 | ): |
| 1008 | """ |
| 1009 | Args: |
| 1010 | binary_mask (ndarray): numpy array of shape (H, W), where H is the image height and |
| 1011 | W is the image width. Each value in the array is either a 0 or 1 value of uint8 |
| 1012 | type. |
| 1013 | color: color of the mask. Refer to `matplotlib.colors` for a full list of |
| 1014 | formats that are accepted. If None, will pick a random color. |
| 1015 | edge_color: color of the polygon edges. Refer to `matplotlib.colors` for a |
| 1016 | full list of formats that are accepted. |
| 1017 | text (str): if None, will be drawn in the object's center of mass. |
| 1018 | alpha (float): blending efficient. Smaller values lead to more transparent masks. |
| 1019 | area_threshold (float): a connected component small than this will not be shown. |
| 1020 | |
| 1021 | Returns: |
| 1022 | output (VisImage): image object with mask drawn. |
| 1023 | """ |
| 1024 | if color is None: |
| 1025 | color = random_color(rgb=True, maximum=1) |
| 1026 | color = mplc.to_rgb(color) |
| 1027 | |
| 1028 | has_valid_segment = False |
| 1029 | binary_mask = binary_mask.astype("uint8") # opencv needs uint8 |
| 1030 | mask = GenericMask(binary_mask, self.output.height, self.output.width) |
| 1031 | shape2d = (binary_mask.shape[0], binary_mask.shape[1]) |
| 1032 | |
| 1033 | if not mask.has_holes: |
| 1034 | # draw polygons for regular masks |
| 1035 | for segment in mask.polygons: |
| 1036 | area = mask_util.area(mask_util.frPyObjects([segment], shape2d[0], shape2d[1])) |
| 1037 | if area < (area_threshold or 0): |
| 1038 | continue |
| 1039 | has_valid_segment = True |
| 1040 | segment = segment.reshape(-1, 2) |
| 1041 | self.draw_polygon(segment, color=color, edge_color=edge_color, alpha=alpha) |
| 1042 | else: |
| 1043 | # TODO: Use Path/PathPatch to draw vector graphics: |
| 1044 | # https://stackoverflow.com/questions/8919719/how-to-plot-a-complex-polygon |
| 1045 | rgba = np.zeros(shape2d + (4,), dtype="float32") |
| 1046 | rgba[:, :, :3] = color |
| 1047 | rgba[:, :, 3] = (mask.mask == 1).astype("float32") * alpha |
| 1048 | has_valid_segment = True |
| 1049 | self.output.ax.imshow(rgba, extent=(0, self.output.width, self.output.height, 0)) |
| 1050 | |
| 1051 | if text is not None and has_valid_segment: |
| 1052 | # TODO sometimes drawn on wrong objects. the heuristics here can improve. |
| 1053 | lighter_color = self._change_color_brightness(color, brightness_factor=0.7) |
| 1054 | _num_cc, cc_labels, stats, centroids = cv2.connectedComponentsWithStats(binary_mask, 8) |
| 1055 | largest_component_id = np.argmax(stats[1:, -1]) + 1 |
| 1056 | |
| 1057 | # draw text on the largest component, as well as other very large components. |
| 1058 | for cid in range(1, _num_cc): |
| 1059 | if cid == largest_component_id or stats[cid, -1] > _LARGE_MASK_AREA_THRESH: |
| 1060 | # median is more stable than centroid |
| 1061 | # center = centroids[largest_component_id] |
| 1062 | center = np.median((cc_labels == cid).nonzero(), axis=1)[::-1] |