# -*- 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