# -*- coding: utf-8 -*- # Copyright (c) Alibaba, Inc. and its affiliates. import copy import io import os import torch import numpy as np import cv2 import imageio from PIL import Image 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() def get_mask_box(mask, threshold=255): locs = np.where(mask >= threshold) if len(locs) < 1 or locs[0].shape[0] < 1 or locs[1].shape[0] < 1: return None left, right = np.min(locs[1]), np.max(locs[1]) top, bottom = np.min(locs[0]), np.max(locs[0]) return [left, top, right, bottom] def convert_to_numpy(image): if isinstance(image, Image.Image): image = np.array(image) elif isinstance(image, torch.Tensor): image = image.detach().cpu().numpy() elif isinstance(image, np.ndarray): image = image.copy() else: raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.' return image def convert_to_pil(image): if isinstance(image, Image.Image): image = image.copy() elif isinstance(image, torch.Tensor): image = image.detach().cpu().numpy() image = Image.fromarray(image.astype('uint8')) elif isinstance(image, np.ndarray): image = Image.fromarray(image.astype('uint8')) else: raise TypeError(f'Unsupported data type {type(image)}, only supports np.ndarray, torch.Tensor, Pillow Image.') return image def convert_to_torch(image): if isinstance(image, Image.Image): image = torch.from_numpy(np.array(image)).float() elif isinstance(image, torch.Tensor): image = image.clone() elif isinstance(image, np.ndarray): image = torch.from_numpy(image.copy()).float() else: raise f'Unsurpport datatype{type(image)}, only surpport np.ndarray, torch.Tensor, Pillow Image.' return image def resize_image(input_image, resolution): H, W, C = input_image.shape H = float(H) W = float(W) k = float(resolution) / min(H, W) H *= k W *= k H = int(np.round(H / 64.0)) * 64 W = int(np.round(W / 64.0)) * 64 img = cv2.resize( input_image, (W, H), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA) return img, k def resize_image_ori(h, w, image, k): img = cv2.resize( image, (w, h), interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA) return img def save_one_video(file_path, videos, fps=8, quality=8, macro_block_size=None): try: video_writer = imageio.get_writer(file_path, fps=fps, codec='libx264', quality=quality, macro_block_size=macro_block_size) for frame in videos: video_writer.append_data(frame) video_writer.close() return True except Exception as e: print(f"Video save error: {e}") return False def save_one_image(file_path, image, use_type='cv2'): try: if use_type == 'cv2': cv2.imwrite(file_path, image) elif use_type == 'pil': image = Image.fromarray(image) image.save(file_path) else: raise ValueError(f"Unknown image write type '{use_type}'") return True except Exception as e: print(f"Image save error: {e}") return False def read_image(image_path, use_type='cv2', is_rgb=True, info=False): image = None width, height = None, None if use_type == 'cv2': try: image = cv2.imread(image_path) if image is None: raise Exception("Image not found or path is incorrect.") if is_rgb: image = image[..., ::-1] height, width = image.shape[:2] except Exception as e: print(f"OpenCV read error: {e}") return None elif use_type == 'pil': try: image = Image.open(image_path) if is_rgb: image = image.convert('RGB') width, height = image.size image = np.array(image) except Exception as e: print(f"PIL read error: {e}") return None else: raise ValueError(f"Unknown image read type '{use_type}'") if info: return image, width, height else: return image def read_mask(mask_path, use_type='cv2', info=False): mask = None width, height = None, None if use_type == 'cv2': try: mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE) if mask is None: raise Exception("Mask not found or path is incorrect.") height, width = mask.shape except Exception as e: print(f"OpenCV read error: {e}") return None elif use_type == 'pil': try: mask = Image.open(mask_path).convert('L') width, height = mask.size mask = np.array(mask) except Exception as e: print(f"PIL read error: {e}") return None else: raise ValueError(f"Unknown mask read type '{use_type}'") if info: return mask, width, height else: return mask def read_video_frames(video_path, use_type='cv2', is_rgb=True, info=False): frames = [] if use_type == "decord": import decord decord.bridge.set_bridge("native") try: cap = decord.VideoReader(video_path) total_frames = len(cap) fps = cap.get_avg_fps() height, width, _ = cap[0].shape frames = [cap[i].asnumpy() for i in range(len(cap))] except Exception as e: print(f"Decord read error: {e}") return None elif use_type == "cv2": try: cap = cv2.VideoCapture(video_path) fps = cap.get(cv2.CAP_PROP_FPS) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) while cap.isOpened(): ret, frame = cap.read() if not ret: break if is_rgb: frames.append(frame[..., ::-1]) else: frames.append(frame) cap.release() total_frames = len(frames) except Exception as e: print(f"OpenCV read error: {e}") return None else: raise ValueError(f"Unknown video type {use_type}") if info: return frames, fps, width, height, total_frames else: return frames def read_video_one_frame(video_path, use_type='cv2', is_rgb=True): image_first = None if use_type == "decord": import decord decord.bridge.set_bridge("native") try: cap = decord.VideoReader(video_path) image_first = cap[0].asnumpy() except Exception as e: print(f"Decord read error: {e}") return None elif use_type == "cv2": try: cap = cv2.VideoCapture(video_path) ret, frame = cap.read() if is_rgb: image_first = frame[..., ::-1] else: image_first = frame cap.release() except Exception as e: print(f"OpenCV read error: {e}") return None else: raise ValueError(f"Unknown video type {use_type}") return image_first def align_frames(first_frame, last_frame): h1, w1 = first_frame.shape[:2] h2, w2 = last_frame.shape[:2] if (h1, w1) == (h2, w2): return last_frame ratio = min(w1 / w2, h1 / h2) new_w = int(w2 * ratio) new_h = int(h2 * ratio) resized = cv2.resize(last_frame, (new_w, new_h), interpolation=cv2.INTER_AREA) aligned = np.ones((h1, w1, 3), dtype=np.uint8) * 255 x_offset = (w1 - new_w) // 2 y_offset = (h1 - new_h) // 2 aligned[y_offset:y_offset + new_h, x_offset:x_offset + new_w] = resized return aligned def save_sam2_video(video_path, video_segments, output_video_path): cap = cv2.VideoCapture(video_path) fps = cap.get(cv2.CAP_PROP_FPS) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) frames = [] while True: ret, frame = cap.read() if not ret: break frames.append(frame) cap.release() obj_mask_map = {} for frame_idx, segments in video_segments.items(): for obj_id, info in segments.items(): seg = single_rle_to_mask(info['mask'])[None, ...].squeeze(0).astype(bool) if obj_id not in obj_mask_map: obj_mask_map[obj_id] = [seg] else: obj_mask_map[obj_id].append(seg) for obj_id, segs in obj_mask_map.items(): output_obj_video_path = os.path.join(output_video_path, f"{obj_id}.mp4") fourcc = cv2.VideoWriter_fourcc(*'mp4v') # codec for saving the video video_writer = cv2.VideoWriter(output_obj_video_path, fourcc, fps, (width * 2, height)) for i, (frame, seg) in enumerate(zip(frames, segs)): print(obj_id, i, np.sum(seg), seg.shape) left_frame = frame.copy() left_frame[seg] = 0 right_frame = frame.copy() right_frame[~seg] = 255 frame_new = np.concatenate([left_frame, right_frame], axis=1) video_writer.write(frame_new) video_writer.release() def get_annotator_instance(anno_cfg): import vace.annotators as annotators anno_cfg = copy.deepcopy(anno_cfg) class_name = anno_cfg.pop("NAME") input_params = anno_cfg.pop("INPUTS") output_params = anno_cfg.pop("OUTPUTS") anno_ins = getattr(annotators, class_name)(cfg=anno_cfg) return {"inputs": input_params, "outputs": output_params, "anno_ins": anno_ins} def get_annotator(config_type='', config_task='', return_dict=True): anno_dict = None from vace.configs import VACE_CONFIGS if config_type in VACE_CONFIGS: task_configs = VACE_CONFIGS[config_type] if config_task in task_configs: anno_dict = get_annotator_instance(task_configs[config_task]) else: raise ValueError(f"Unknown config task {config_task}") else: for cfg_type, cfg_dict in VACE_CONFIGS.items(): if config_task in cfg_dict: for task_name, task_cfg in cfg_dict[config_task].items(): anno_dict = get_annotator_instance(task_cfg) else: raise ValueError(f"Unknown config type {config_type}") if return_dict: return anno_dict else: return anno_dict['anno_ins']