61 lines
2.2 KiB
Python
61 lines
2.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
import cv2
|
|
import numpy as np
|
|
import torch
|
|
|
|
from .utils import convert_to_numpy
|
|
|
|
|
|
class SubjectAnnotator:
|
|
def __init__(self, cfg, device=None):
|
|
self.mode = cfg.get('MODE', "salientmasktrack")
|
|
self.use_aug = cfg.get('USE_AUG', False)
|
|
self.use_crop = cfg.get('USE_CROP', False)
|
|
self.roi_only = cfg.get('ROI_ONLY', False)
|
|
self.return_mask = cfg.get('RETURN_MASK', True)
|
|
|
|
from .inpainting import InpaintingAnnotator
|
|
self.inp_anno = InpaintingAnnotator(cfg['INPAINTING'], device=device)
|
|
if self.use_aug:
|
|
from .maskaug import MaskAugAnnotator
|
|
self.maskaug_anno = MaskAugAnnotator(cfg={})
|
|
assert self.mode in ["plain", "salient", "mask", "bbox", "salientmasktrack", "salientbboxtrack", "masktrack",
|
|
"bboxtrack", "label", "caption", "all"]
|
|
|
|
def forward(self, image=None, mode=None, return_mask=None, mask_cfg=None, mask=None, bbox=None, label=None, caption=None):
|
|
return_mask = return_mask if return_mask is not None else self.return_mask
|
|
|
|
if mode == "plain":
|
|
return {"image": image, "mask": None} if return_mask else image
|
|
|
|
inp_res = self.inp_anno.forward(image, mask=mask, bbox=bbox, label=label, caption=caption, mode=mode, return_mask=True, return_source=True)
|
|
src_image = inp_res['src_image']
|
|
mask = inp_res['mask']
|
|
|
|
if self.use_aug and mask_cfg is not None:
|
|
mask = self.maskaug_anno.forward(mask, mask_cfg)
|
|
|
|
_, binary_mask = cv2.threshold(mask, 1, 255, cv2.THRESH_BINARY)
|
|
if (binary_mask is None or binary_mask.size == 0 or cv2.countNonZero(binary_mask) == 0):
|
|
x, y, w, h = 0, 0, binary_mask.shape[1], binary_mask.shape[0]
|
|
else:
|
|
x, y, w, h = cv2.boundingRect(binary_mask)
|
|
|
|
ret_mask = mask.copy()
|
|
ret_image = src_image.copy()
|
|
|
|
if self.roi_only:
|
|
ret_image[mask == 0] = 255
|
|
|
|
if self.use_crop:
|
|
ret_image = ret_image[y:y + h, x:x + w]
|
|
ret_mask = ret_mask[y:y + h, x:x + w]
|
|
|
|
if return_mask:
|
|
return {"image": ret_image, "mask": ret_mask}
|
|
else:
|
|
return ret_image
|
|
|
|
|