|
|
|
import numpy as np |
|
from typing import List, Union, Tuple |
|
import torch |
|
from utils.constants import COLOR_PALETTE |
|
from utils.constants import get_color |
|
import cv2 |
|
|
|
def tags2multilines(tags: Union[str, List], lw, tf, max_width): |
|
if isinstance(tags, str): |
|
taglist = tags.split(' ') |
|
else: |
|
taglist = tags |
|
|
|
sz = cv2.getTextSize(' ', 0, lw / 3, tf) |
|
line_height = sz[0][1] |
|
line_width = 0 |
|
if len(taglist) > 0: |
|
lines = [taglist[0]] |
|
if len(taglist) > 1: |
|
for t in taglist[1:]: |
|
textl = len(t) * line_height |
|
if line_width + line_height + textl > max_width: |
|
lines.append(t) |
|
line_width = 0 |
|
else: |
|
line_width = line_width + line_height + textl |
|
lines[-1] = lines[-1] + ' ' + t |
|
return lines, line_height |
|
|
|
class AnimeInstances: |
|
|
|
def __init__(self, |
|
masks: Union[np.ndarray, torch.Tensor ]= None, |
|
bboxes: Union[np.ndarray, torch.Tensor ] = None, |
|
scores: Union[np.ndarray, torch.Tensor ] = None, |
|
tags: List[str] = None, character_tags: List[str] = None) -> None: |
|
self.masks = masks |
|
self.tags = tags |
|
self.bboxes = bboxes |
|
|
|
|
|
if scores is None: |
|
scores = [1.] * len(self) |
|
if self.is_numpy: |
|
scores = np.array(scores) |
|
elif self.is_tensor: |
|
scores = torch.tensor(scores) |
|
|
|
self.scores = scores |
|
|
|
if tags is None: |
|
self.tags = [''] * len(self) |
|
self.character_tags = [''] * len(self) |
|
else: |
|
self.tags = tags |
|
self.character_tags = character_tags |
|
|
|
@property |
|
def is_cuda(self): |
|
if isinstance(self.masks, torch.Tensor) and self.masks.is_cuda: |
|
return True |
|
else: |
|
return False |
|
|
|
@property |
|
def is_tensor(self): |
|
if self.is_empty: |
|
return False |
|
else: |
|
return isinstance(self.masks, torch.Tensor) |
|
|
|
@property |
|
def is_numpy(self): |
|
if self.is_empty: |
|
return True |
|
else: |
|
return isinstance(self.masks, np.ndarray) |
|
|
|
@property |
|
def is_empty(self): |
|
return self.masks is None or len(self.masks) == 0\ |
|
|
|
def remove_duplicated(self): |
|
|
|
num_masks = len(self) |
|
if num_masks < 2: |
|
return |
|
|
|
need_cvt = False |
|
if self.is_numpy: |
|
need_cvt = True |
|
self.to_tensor() |
|
|
|
mask_areas = torch.Tensor([mask.sum() for mask in self.masks]) |
|
sids = torch.argsort(mask_areas, descending=True) |
|
sids = sids.cpu().numpy().tolist() |
|
mask_areas = mask_areas[sids] |
|
masks = self.masks[sids] |
|
bboxes = self.bboxes[sids] |
|
tags = [self.tags[sid] for sid in sids] |
|
scores = self.scores[sids] |
|
|
|
canvas = masks[0] |
|
|
|
valid_ids: List = np.arange(num_masks).tolist() |
|
for ii, mask in enumerate(masks[1:]): |
|
|
|
mask_id = ii + 1 |
|
canvas_and = torch.bitwise_and(canvas, mask) |
|
|
|
and_area = canvas_and.sum() |
|
mask_area = mask_areas[mask_id] |
|
|
|
if and_area / mask_area > 0.8: |
|
valid_ids.remove(mask_id) |
|
elif mask_id != num_masks - 1: |
|
canvas = torch.bitwise_or(canvas, mask) |
|
|
|
sids = valid_ids |
|
self.masks = masks[sids] |
|
self.bboxes = bboxes[sids] |
|
self.tags = [tags[sid] for sid in sids] |
|
self.scores = scores[sids] |
|
|
|
if need_cvt: |
|
self.to_numpy() |
|
|
|
|
|
|
|
def draw_instances(self, |
|
img: np.ndarray, |
|
draw_bbox: bool = True, |
|
draw_ins_mask: bool = True, |
|
draw_ins_contour: bool = True, |
|
draw_tags: bool = False, |
|
draw_indices: List = None, |
|
mask_alpha: float = 0.4): |
|
|
|
mask_alpha = 0.75 |
|
|
|
|
|
drawed = img.copy() |
|
|
|
if self.is_empty: |
|
return drawed |
|
|
|
im_h, im_w = img.shape[:2] |
|
|
|
mask_shape = self.masks[0].shape |
|
if mask_shape[0] != im_h or mask_shape[1] != im_w: |
|
drawed = cv2.resize(drawed, (mask_shape[1], mask_shape[0]), interpolation=cv2.INTER_AREA) |
|
im_h, im_w = mask_shape[0], mask_shape[1] |
|
|
|
if draw_indices is None: |
|
draw_indices = list(range(len(self))) |
|
ins_dict = {'mask': [], 'tags': [], 'score': [], 'bbox': [], 'character_tags': []} |
|
colors = [] |
|
for idx in draw_indices: |
|
ins = self.get_instance(idx, out_type='numpy') |
|
for key, data in ins.items(): |
|
ins_dict[key].append(data) |
|
colors.append(get_color(idx)) |
|
|
|
if draw_bbox: |
|
lw = max(round(sum(drawed.shape) / 2 * 0.003), 2) |
|
for color, bbox in zip(colors, ins_dict['bbox']): |
|
p1, p2 = (int(bbox[0]), int(bbox[1])), (int(bbox[2] + bbox[0]), int(bbox[3] + bbox[1])) |
|
cv2.rectangle(drawed, p1, p2, color, thickness=lw, lineType=cv2.LINE_AA) |
|
|
|
if draw_ins_mask: |
|
drawed = drawed.astype(np.float32) |
|
for color, mask in zip(colors, ins_dict['mask']): |
|
p = mask.astype(np.float32) |
|
blend_mask = np.full((im_h, im_w, 3), color, dtype=np.float32) |
|
alpha_msk = (mask_alpha * p)[..., None] |
|
alpha_ori = 1 - alpha_msk |
|
drawed = drawed * alpha_ori + alpha_msk * blend_mask |
|
drawed = drawed.astype(np.uint8) |
|
|
|
if draw_tags: |
|
lw = max(round(sum(drawed.shape) / 2 * 0.002), 2) |
|
tf = max(lw - 1, 1) |
|
for color, tags, bbox in zip(colors, ins_dict['tags'], ins_dict['bbox']): |
|
if not tags: |
|
continue |
|
lines, line_height = tags2multilines(tags, lw, tf, bbox[2]) |
|
for ii, l in enumerate(lines): |
|
xy = (bbox[0], bbox[1] + line_height + int(line_height * 1.2 * ii)) |
|
cv2.putText(drawed, l, xy, 0, lw / 3, color, thickness=tf, lineType=cv2.LINE_AA) |
|
|
|
|
|
|
|
return drawed |
|
|
|
|
|
def cuda(self): |
|
if self.is_empty: |
|
return self |
|
self.to_tensor(device='cuda') |
|
return self |
|
|
|
def cpu(self): |
|
if not self.is_tensor or not self.is_cuda: |
|
return self |
|
self.masks = self.masks.cpu() |
|
self.scores = self.scores.cpu() |
|
self.bboxes = self.bboxes.cpu() |
|
return self |
|
|
|
def to_tensor(self, device: str = 'cpu'): |
|
if self.is_empty: |
|
return self |
|
elif self.is_tensor and self.masks.device == device: |
|
return self |
|
self.masks = torch.from_numpy(self.masks).to(device) |
|
self.bboxes = torch.from_numpy(self.bboxes).to(device) |
|
self.scores = torch.from_numpy(self.scores ).to(device) |
|
return self |
|
|
|
def to_numpy(self): |
|
if self.is_numpy: |
|
return self |
|
if self.is_cuda: |
|
self.masks = self.masks.cpu().numpy() |
|
self.scores = self.scores.cpu().numpy() |
|
self.bboxes = self.bboxes.cpu().numpy() |
|
else: |
|
self.masks = self.masks.numpy() |
|
self.scores = self.scores.numpy() |
|
self.bboxes = self.bboxes.numpy() |
|
return self |
|
|
|
def get_instance(self, ins_idx: int, out_type: str = None, device: str = None): |
|
mask = self.masks[ins_idx] |
|
tags = self.tags[ins_idx] |
|
character_tags = self.character_tags[ins_idx] |
|
bbox = self.bboxes[ins_idx] |
|
score = self.scores[ins_idx] |
|
if out_type is not None: |
|
if out_type == 'numpy' and not self.is_numpy: |
|
mask = mask.cpu().numpy() |
|
bbox = bbox.cpu().numpy() |
|
score = score.cpu().numpy() |
|
if out_type == 'tensor' and not self.is_tensor: |
|
mask = torch.from_numpy(mask) |
|
bbox = torch.from_numpy(bbox) |
|
score = torch.from_numpy(score) |
|
if isinstance(mask, torch.Tensor) and device is not None and mask.device != device: |
|
mask = mask.to(device) |
|
bbox = bbox.to(device) |
|
score = score.to(device) |
|
|
|
return { |
|
'mask': mask, |
|
'tags': tags, |
|
'character_tags': character_tags, |
|
'bbox': bbox, |
|
'score': score |
|
} |
|
|
|
def __len__(self): |
|
if self.is_empty: |
|
return 0 |
|
else: |
|
return len(self.masks) |
|
|
|
def resize(self, h, w, mode = 'area'): |
|
if self.is_empty: |
|
return |
|
if self.is_tensor: |
|
masks = self.masks.to(torch.float).unsqueeze(1) |
|
oh, ow = masks.shape[2], masks.shape[3] |
|
hs, ws = h / oh, w / ow |
|
bboxes = self.bboxes.float() |
|
bboxes[:, ::2] *= hs |
|
bboxes[:, 1::2] *= ws |
|
self.bboxes = torch.round(bboxes).int() |
|
masks = torch.nn.functional.interpolate(masks, (h, w), mode=mode) |
|
self.masks = masks.squeeze(1) > 0.3 |
|
|
|
def compose_masks(self, output_type=None): |
|
if self.is_empty: |
|
return None |
|
else: |
|
mask = self.masks[0] |
|
if len(self.masks) > 1: |
|
for m in self.masks[1:]: |
|
if self.is_numpy: |
|
mask = np.logical_or(mask, m) |
|
else: |
|
mask = torch.logical_or(mask, m) |
|
if output_type is not None: |
|
if output_type == 'numpy' and not self.is_numpy: |
|
mask = mask.cpu().numpy() |
|
if output_type == 'tensor' and not self.is_tensor: |
|
mask = torch.from_numpy(mask) |
|
return mask |
|
|
|
|
|
|
|
|