162 lines
6.8 KiB
Python
162 lines
6.8 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from .utils import convert_to_numpy
|
|
|
|
|
|
class LayoutBboxAnnotator:
|
|
def __init__(self, cfg, device=None):
|
|
self.bg_color = cfg.get('BG_COLOR', [255, 255, 255])
|
|
self.box_color = cfg.get('BOX_COLOR', [0, 0, 0])
|
|
self.frame_size = cfg.get('FRAME_SIZE', [720, 1280]) # [H, W]
|
|
self.num_frames = cfg.get('NUM_FRAMES', 81)
|
|
ram_tag_color_path = cfg.get('RAM_TAG_COLOR_PATH', None)
|
|
self.color_dict = {'default': tuple(self.box_color)}
|
|
if ram_tag_color_path is not None:
|
|
lines = [id_name_color.strip().split('#;#') for id_name_color in open(ram_tag_color_path).readlines()]
|
|
self.color_dict.update({id_name_color[1]: tuple(eval(id_name_color[2])) for id_name_color in lines})
|
|
|
|
def forward(self, bbox, frame_size=None, num_frames=None, label=None, color=None):
|
|
frame_size = frame_size if frame_size is not None else self.frame_size
|
|
num_frames = num_frames if num_frames is not None else self.num_frames
|
|
assert len(bbox) == 2, 'bbox should be a list of two elements (start_bbox & end_bbox)'
|
|
# frame_size = [H, W]
|
|
# bbox = [x1, y1, x2, y2]
|
|
label = label[0] if label is not None and isinstance(label, list) else label
|
|
if label is not None and label in self.color_dict:
|
|
box_color = self.color_dict[label]
|
|
elif color is not None:
|
|
box_color = color
|
|
else:
|
|
box_color = self.color_dict['default']
|
|
start_bbox, end_bbox = bbox
|
|
start_bbox = [start_bbox[0], start_bbox[1], start_bbox[2] - start_bbox[0], start_bbox[3] - start_bbox[1]]
|
|
start_bbox = np.array(start_bbox, dtype=np.float32)
|
|
end_bbox = [end_bbox[0], end_bbox[1], end_bbox[2] - end_bbox[0], end_bbox[3] - end_bbox[1]]
|
|
end_bbox = np.array(end_bbox, dtype=np.float32)
|
|
bbox_increment = (end_bbox - start_bbox) / num_frames
|
|
ret_frames = []
|
|
for frame_idx in range(num_frames):
|
|
frame = np.zeros((frame_size[0], frame_size[1], 3), dtype=np.uint8)
|
|
frame[:] = self.bg_color
|
|
current_bbox = start_bbox + bbox_increment * frame_idx
|
|
current_bbox = current_bbox.astype(int)
|
|
x, y, w, h = current_bbox
|
|
cv2.rectangle(frame, (x, y), (x + w, y + h), box_color, 2)
|
|
ret_frames.append(frame[..., ::-1])
|
|
return ret_frames
|
|
|
|
|
|
|
|
|
|
class LayoutMaskAnnotator:
|
|
def __init__(self, cfg, device=None):
|
|
self.use_aug = cfg.get('USE_AUG', False)
|
|
self.bg_color = cfg.get('BG_COLOR', [255, 255, 255])
|
|
self.box_color = cfg.get('BOX_COLOR', [0, 0, 0])
|
|
ram_tag_color_path = cfg.get('RAM_TAG_COLOR_PATH', None)
|
|
self.color_dict = {'default': tuple(self.box_color)}
|
|
if ram_tag_color_path is not None:
|
|
lines = [id_name_color.strip().split('#;#') for id_name_color in open(ram_tag_color_path).readlines()]
|
|
self.color_dict.update({id_name_color[1]: tuple(eval(id_name_color[2])) for id_name_color in lines})
|
|
if self.use_aug:
|
|
from .maskaug import MaskAugAnnotator
|
|
self.maskaug_anno = MaskAugAnnotator(cfg={})
|
|
|
|
|
|
def find_contours(self, mask):
|
|
contours, hier = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
return contours
|
|
|
|
def draw_contours(self, canvas, contour, color):
|
|
canvas = np.ascontiguousarray(canvas, dtype=np.uint8)
|
|
canvas = cv2.drawContours(canvas, contour, -1, color, thickness=3)
|
|
return canvas
|
|
|
|
def forward(self, mask=None, color=None, label=None, mask_cfg=None):
|
|
if not isinstance(mask, list):
|
|
is_batch = False
|
|
mask = [mask]
|
|
else:
|
|
is_batch = True
|
|
|
|
if label is not None and label in self.color_dict:
|
|
color = self.color_dict[label]
|
|
elif color is not None:
|
|
color = color
|
|
else:
|
|
color = self.color_dict['default']
|
|
|
|
ret_data = []
|
|
for sub_mask in mask:
|
|
sub_mask = convert_to_numpy(sub_mask)
|
|
if self.use_aug:
|
|
sub_mask = self.maskaug_anno.forward(sub_mask, mask_cfg)
|
|
canvas = np.ones((sub_mask.shape[0], sub_mask.shape[1], 3)) * 255
|
|
contour = self.find_contours(sub_mask)
|
|
frame = self.draw_contours(canvas, contour, color)
|
|
ret_data.append(frame)
|
|
|
|
if is_batch:
|
|
return ret_data
|
|
else:
|
|
return ret_data[0]
|
|
|
|
|
|
|
|
|
|
class LayoutTrackAnnotator:
|
|
def __init__(self, cfg, device=None):
|
|
self.use_aug = cfg.get('USE_AUG', False)
|
|
self.bg_color = cfg.get('BG_COLOR', [255, 255, 255])
|
|
self.box_color = cfg.get('BOX_COLOR', [0, 0, 0])
|
|
ram_tag_color_path = cfg.get('RAM_TAG_COLOR_PATH', None)
|
|
self.color_dict = {'default': tuple(self.box_color)}
|
|
if ram_tag_color_path is not None:
|
|
lines = [id_name_color.strip().split('#;#') for id_name_color in open(ram_tag_color_path).readlines()]
|
|
self.color_dict.update({id_name_color[1]: tuple(eval(id_name_color[2])) for id_name_color in lines})
|
|
if self.use_aug:
|
|
from .maskaug import MaskAugAnnotator
|
|
self.maskaug_anno = MaskAugAnnotator(cfg={})
|
|
from .inpainting import InpaintingVideoAnnotator
|
|
self.inpainting_anno = InpaintingVideoAnnotator(cfg=cfg['INPAINTING'])
|
|
|
|
def find_contours(self, mask):
|
|
contours, hier = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
return contours
|
|
|
|
def draw_contours(self, canvas, contour, color):
|
|
canvas = np.ascontiguousarray(canvas, dtype=np.uint8)
|
|
canvas = cv2.drawContours(canvas, contour, -1, color, thickness=3)
|
|
return canvas
|
|
|
|
def forward(self, color=None, mask_cfg=None, frames=None, video=None, mask=None, bbox=None, label=None, caption=None, mode=None):
|
|
inp_data = self.inpainting_anno.forward(frames, video, mask, bbox, label, caption, mode)
|
|
inp_masks = inp_data['masks']
|
|
|
|
label = label[0] if label is not None and isinstance(label, list) else label
|
|
if label is not None and label in self.color_dict:
|
|
color = self.color_dict[label]
|
|
elif color is not None:
|
|
color = color
|
|
else:
|
|
color = self.color_dict['default']
|
|
|
|
num_frames = len(inp_masks)
|
|
ret_data = []
|
|
for i in range(num_frames):
|
|
sub_mask = inp_masks[i]
|
|
if self.use_aug and mask_cfg is not None:
|
|
sub_mask = self.maskaug_anno.forward(sub_mask, mask_cfg)
|
|
canvas = np.ones((sub_mask.shape[0], sub_mask.shape[1], 3)) * 255
|
|
contour = self.find_contours(sub_mask)
|
|
frame = self.draw_contours(canvas, contour, color)
|
|
ret_data.append(frame)
|
|
|
|
return ret_data
|
|
|
|
|