|
|
|
|
|
|
|
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]) |
|
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)' |
|
|
|
|
|
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 |
|
|
|
|
|
|
