154 lines
6.0 KiB
Python
154 lines
6.0 KiB
Python
# -*- coding: utf-8 -*-
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
from abc import ABCMeta
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
from scipy import ndimage
|
|
try:
|
|
from sklearn.cluster import KMeans
|
|
except:
|
|
import warnings
|
|
warnings.warn("ignore sklearn import, please pip install scikit-learn.")
|
|
|
|
from scepter.modules.annotator.base_annotator import BaseAnnotator
|
|
from scepter.modules.annotator.registry import ANNOTATORS
|
|
from scepter.modules.utils.config import dict_to_yaml
|
|
from scepter.modules.utils.file_system import FS
|
|
import pycocotools.mask as mask_utils
|
|
|
|
|
|
def single_mask_to_rle(mask):
|
|
rle = mask_utils.encode(np.array(mask[:, :, None], order="F", dtype="uint8"))[0]
|
|
rle["counts"] = rle["counts"].decode("utf-8")
|
|
return rle
|
|
|
|
def single_rle_to_mask(rle):
|
|
mask = np.array(mask_utils.decode(rle)).astype(np.uint8)
|
|
return mask
|
|
|
|
def single_mask_to_xyxy(mask):
|
|
bbox = np.zeros((4), dtype=int)
|
|
rows, cols = np.where(np.array(mask))
|
|
if len(rows) > 0 and len(cols) > 0:
|
|
x_min, x_max = np.min(cols), np.max(cols)
|
|
y_min, y_max = np.min(rows), np.max(rows)
|
|
bbox[:] = [x_min, y_min, x_max, y_max]
|
|
return bbox.tolist()
|
|
|
|
@ANNOTATORS.register_class()
|
|
class SAM2DrawVideoAnnotator(BaseAnnotator, metaclass=ABCMeta):
|
|
para_dict = {}
|
|
|
|
def __init__(self, cfg, logger=None):
|
|
super().__init__(cfg, logger=logger)
|
|
self.task_type = cfg.get('TASK_TYPE', 'input_box')
|
|
from sam2.build_sam import build_sam2_video_predictor
|
|
config_path = FS.get_from(cfg.CONFIG_PATH, local_path=cfg.CONFIG_LOCAL_PATH, wait_finish=True)
|
|
pretrained_model = FS.get_from(cfg.PRETRAINED_MODEL, wait_finish=True)
|
|
self.video_predictor = build_sam2_video_predictor(config_path, pretrained_model, fill_hole_area=0)
|
|
|
|
def forward(self,
|
|
video,
|
|
input_box=None,
|
|
mask=None,
|
|
task_type=None):
|
|
task_type = task_type if task_type is not None else self.task_type
|
|
|
|
if mask is not None:
|
|
if isinstance(mask, Image.Image):
|
|
mask = np.array(mask)
|
|
elif isinstance(mask, torch.Tensor):
|
|
mask = mask.detach().cpu().numpy()
|
|
elif isinstance(mask, np.ndarray):
|
|
mask = mask.copy()
|
|
else:
|
|
raise f'Unsurpport datatype{type(mask)}, only surpport np.ndarray, torch.Tensor, Pillow Image.'
|
|
|
|
if task_type == 'mask_point':
|
|
if len(mask.shape) == 3:
|
|
scribble = mask.transpose(2, 1, 0)[0]
|
|
else:
|
|
scribble = mask.transpose(1, 0) # (H, W) -> (W, H)
|
|
labeled_array, num_features = ndimage.label(scribble >= 255)
|
|
centers = ndimage.center_of_mass(scribble, labeled_array,
|
|
range(1, num_features + 1))
|
|
point_coords = np.array(centers)
|
|
point_labels = np.array([1] * len(centers))
|
|
sample = {
|
|
'points': point_coords,
|
|
'labels': point_labels
|
|
}
|
|
elif task_type == 'mask_box':
|
|
if len(mask.shape) == 3:
|
|
scribble = mask.transpose(2, 1, 0)[0]
|
|
else:
|
|
scribble = mask.transpose(1, 0) # (H, W) -> (W, H)
|
|
labeled_array, num_features = ndimage.label(scribble >= 255)
|
|
centers = ndimage.center_of_mass(scribble, labeled_array,
|
|
range(1, num_features + 1))
|
|
centers = np.array(centers)
|
|
# (x1, y1, x2, y2)
|
|
x_min = centers[:, 0].min()
|
|
x_max = centers[:, 0].max()
|
|
y_min = centers[:, 1].min()
|
|
y_max = centers[:, 1].max()
|
|
bbox = np.array([x_min, y_min, x_max, y_max])
|
|
sample = {'box': bbox}
|
|
elif task_type == 'input_box':
|
|
if isinstance(input_box, list):
|
|
input_box = np.array(input_box)
|
|
sample = {'box': input_box}
|
|
elif task_type == 'mask':
|
|
sample = {'mask': mask}
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
ann_frame_idx = 0
|
|
object_id = 0
|
|
with (torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16)):
|
|
|
|
inference_state = self.video_predictor.init_state(video_path=video)
|
|
|
|
if task_type in ['mask_point', 'mask_box', 'input_box']:
|
|
_, out_obj_ids, out_mask_logits = self.video_predictor.add_new_points_or_box(
|
|
inference_state=inference_state,
|
|
frame_idx=ann_frame_idx,
|
|
obj_id=object_id,
|
|
**sample
|
|
)
|
|
elif task_type in ['mask']:
|
|
_, out_obj_ids, out_mask_logits = self.video_predictor.add_new_mask(
|
|
inference_state=inference_state,
|
|
frame_idx=ann_frame_idx,
|
|
obj_id=object_id,
|
|
**sample
|
|
)
|
|
else:
|
|
raise NotImplementedError
|
|
|
|
video_segments = {} # video_segments contains the per-frame segmentation results
|
|
for out_frame_idx, out_obj_ids, out_mask_logits in self.video_predictor.propagate_in_video(inference_state):
|
|
frame_segments = {}
|
|
for i, out_obj_id in enumerate(out_obj_ids):
|
|
mask = (out_mask_logits[i] > 0.0).cpu().numpy().squeeze(0)
|
|
frame_segments[out_obj_id] = {
|
|
"mask": single_mask_to_rle(mask),
|
|
"mask_area": int(mask.sum()),
|
|
"mask_box": single_mask_to_xyxy(mask),
|
|
}
|
|
video_segments[out_frame_idx] = frame_segments
|
|
|
|
ret_data = {
|
|
"annotations": video_segments
|
|
}
|
|
return ret_data
|
|
|
|
@staticmethod
|
|
def get_config_template():
|
|
return dict_to_yaml('ANNOTATORS',
|
|
__class__.__name__,
|
|
SAM2DrawVideoAnnotator.para_dict,
|
|
set_name=True)
|