Spaces:
Build error
Build error
| import numpy as np | |
| import pycocotools.mask as mask_util | |
| import cv2 | |
| class GenericMask: | |
| """ | |
| Attribute: | |
| polygons (list[ndarray]): list[ndarray]: polygons for this mask. | |
| Each ndarray has format [x, y, x, y, ...] | |
| mask (ndarray): a binary mask | |
| """ | |
| def __init__(self, mask_or_polygons, height, width): | |
| self._mask = self._polygons = self._has_holes = None | |
| self.height = height | |
| self.width = width | |
| m = mask_or_polygons | |
| if isinstance(m, dict): | |
| # RLEs | |
| assert "counts" in m and "size" in m | |
| if isinstance(m["counts"], list): # uncompressed RLEs | |
| h, w = m["size"] | |
| assert h == height and w == width | |
| m = mask_util.frPyObjects(m, h, w) | |
| self._mask = mask_util.decode(m)[:, :] | |
| return | |
| if isinstance(m, list): # list[ndarray] | |
| self._polygons = [np.asarray(x).reshape(-1) for x in m] | |
| return | |
| if isinstance(m, np.ndarray): # assumed to be a binary mask | |
| assert m.shape[1] != 2, m.shape | |
| assert m.shape == (height, width), m.shape | |
| self._mask = m.astype("uint8") | |
| return | |
| raise ValueError("GenericMask cannot handle object {} of type '{}'".format(m, type(m))) | |
| def mask(self): | |
| if self._mask is None: | |
| self._mask = self.polygons_to_mask(self._polygons) | |
| return self._mask | |
| def polygons(self): | |
| if self._polygons is None: | |
| self._polygons, self._has_holes = self.mask_to_polygons(self._mask) | |
| return self._polygons | |
| def has_holes(self): | |
| if self._has_holes is None: | |
| if self._mask is not None: | |
| self._polygons, self._has_holes = self.mask_to_polygons(self._mask) | |
| else: | |
| self._has_holes = False # if original format is polygon, does not have holes | |
| return self._has_holes | |
| def mask_to_polygons(self, mask): | |
| # cv2.RETR_CCOMP flag retrieves all the contours and arranges them to a 2-level | |
| # hierarchy. External contours (boundary) of the object are placed in hierarchy-1. | |
| # Internal contours (holes) are placed in hierarchy-2. | |
| # cv2.CHAIN_APPROX_NONE flag gets vertices of polygons from contours. | |
| mask = np.ascontiguousarray(mask) # some versions of cv2 does not support incontiguous arr | |
| #res = cv2.findContours(mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE) | |
| res = cv2.findContours(mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE) | |
| hierarchy = res[-1] | |
| if hierarchy is None: # empty mask | |
| return [], False | |
| has_holes = (hierarchy.reshape(-1, 4)[:, 3] >= 0).sum() > 0 | |
| res = res[-2] | |
| res = [x.flatten() for x in res] | |
| # These coordinates from OpenCV are integers in range [0, W-1 or H-1]. | |
| # We add 0.5 to turn them into real-value coordinate space. A better solution | |
| # would be to first +0.5 and then dilate the returned polygon by 0.5. | |
| res = [x + 0.5 for x in res if len(x) >= 6] | |
| return res, has_holes | |
| def polygons_to_mask(self, polygons): | |
| rle = mask_util.frPyObjects(polygons, self.height, self.width) | |
| rle = mask_util.merge(rle) | |
| return mask_util.decode(rle)[:, :] | |
| def area(self): | |
| return self.mask.sum() | |
| def bbox(self): | |
| p = mask_util.frPyObjects(self.polygons, self.height, self.width) | |
| p = mask_util.merge(p) | |
| bbox = mask_util.toBbox(p) | |
| bbox[2] += bbox[0] | |
| bbox[3] += bbox[1] | |
| return bbox | |