File size: 3,197 Bytes
690f890 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import numpy as np
from scipy.spatial import ConvexHull
from skimage.draw import polygon
from scipy import ndimage
from .utils import convert_to_numpy
class MaskDrawAnnotator:
def __init__(self, cfg, device=None):
self.mode = cfg.get('MODE', 'maskpoint')
self.return_dict = cfg.get('RETURN_DICT', True)
assert self.mode in ['maskpoint', 'maskbbox', 'mask', 'bbox']
def forward(self,
mask=None,
image=None,
bbox=None,
mode=None,
return_dict=None):
mode = mode if mode is not None else self.mode
return_dict = return_dict if return_dict is not None else self.return_dict
mask = convert_to_numpy(mask) if mask is not None else None
image = convert_to_numpy(image) if image is not None else None
mask_shape = mask.shape
if mode == 'maskpoint':
scribble = mask.transpose(1, 0)
labeled_array, num_features = ndimage.label(scribble >= 255)
centers = ndimage.center_of_mass(scribble, labeled_array,
range(1, num_features + 1))
centers = np.array(centers)
out_mask = np.zeros(mask_shape, dtype=np.uint8)
hull = ConvexHull(centers)
hull_vertices = centers[hull.vertices]
rr, cc = polygon(hull_vertices[:, 1], hull_vertices[:, 0], mask_shape)
out_mask[rr, cc] = 255
elif mode == 'maskbbox':
scribble = mask.transpose(1, 0)
labeled_array, num_features = ndimage.label(scribble >= 255)
centers = ndimage.center_of_mass(scribble, labeled_array,
range(1, num_features + 1))
centers = np.array(centers)
# (x1, y1, x2, y2)
x_min = centers[:, 0].min()
x_max = centers[:, 0].max()
y_min = centers[:, 1].min()
y_max = centers[:, 1].max()
out_mask = np.zeros(mask_shape, dtype=np.uint8)
out_mask[int(y_min) : int(y_max) + 1, int(x_min) : int(x_max) + 1] = 255
if image is not None:
out_image = image[int(y_min) : int(y_max) + 1, int(x_min) : int(x_max) + 1]
elif mode == 'bbox':
if isinstance(bbox, list):
bbox = np.array(bbox)
x_min, y_min, x_max, y_max = bbox
out_mask = np.zeros(mask_shape, dtype=np.uint8)
out_mask[int(y_min) : int(y_max) + 1, int(x_min) : int(x_max) + 1] = 255
if image is not None:
out_image = image[int(y_min) : int(y_max) + 1, int(x_min) : int(x_max) + 1]
elif mode == 'mask':
out_mask = mask
else:
raise NotImplementedError
if return_dict:
if image is not None:
return {"image": out_image, "mask": out_mask}
else:
return {"mask": out_mask}
else:
if image is not None:
return out_image, out_mask
else:
return out_mask |