79 lines
3.1 KiB
Python
79 lines
3.1 KiB
Python
# -*- 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 |