# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved. import os import numpy as np import shutil import torch from diffusers import FluxKontextPipeline import cv2 from loguru import logger from PIL import Image try: import moviepy.editor as mpy except: import moviepy as mpy from decord import VideoReader from pose2d import Pose2d from pose2d_utils import AAPoseMeta from utils import resize_by_area, get_frame_indices, padding_resize, get_face_bboxes, get_aug_mask, get_mask_body_img from human_visualization import draw_aapose_by_meta_new from retarget_pose import get_retarget_pose from sam2.build_sam import build_sam2, build_sam2_video_predictor def get_frames(video_path, resolution_area, fps=30): video_reader = VideoReader(video_path) frame_num = len(video_reader) video_fps = video_reader.get_avg_fps() # TODO: Maybe we can switch to PyAV later, which can get accurate frame num duration = video_reader.get_frame_timestamp(-1)[-1] expected_frame_num = int(duration * video_fps + 0.5) ratio = abs((frame_num - expected_frame_num)/frame_num) if ratio > 0.1: print("Warning: The difference between the actual number of frames and the expected number of frames is two large") frame_num = expected_frame_num if fps == -1: fps = video_fps target_num = int(frame_num / video_fps * fps) idxs = get_frame_indices(frame_num, video_fps, target_num, fps) frames = video_reader.get_batch(idxs).asnumpy() frames = [resize_by_area(frame, resolution_area[0] * resolution_area[1], divisor=16) for frame in frames] return frames def quantize_mask_blocky(mask, block_w=16, block_h=16, occupancy=0.15): """ Convert a binary mask to a blocky (quantized) mask. - block_w, block_h: target block size in pixels - occupancy: fraction [0..1] of foreground within a block to turn it on """ m = (mask > 0).astype(np.uint8) H, W = m.shape[:2] # compute “block grid” size grid_w = max(1, int(np.ceil(W / block_w))) grid_h = max(1, int(np.ceil(H / block_h))) # downsample to grid using area interpolation (captures occupancy) small = cv2.resize(m, (grid_w, grid_h), interpolation=cv2.INTER_AREA) # threshold by occupancy (values now in [0,1] if source was 0/1) small_q = (small >= occupancy).astype(np.uint8) # upsample back with nearest (keeps sharp blocks) blocky = cv2.resize(small_q, (W, H), interpolation=cv2.INTER_NEAREST) return blocky class ProcessPipeline(): def __init__(self, det_checkpoint_path, pose2d_checkpoint_path, sam_checkpoint_path, flux_kontext_path): self.pose2d = Pose2d(checkpoint=pose2d_checkpoint_path, detector_checkpoint=det_checkpoint_path) model_cfg = "sam2_hiera_l.yaml" if sam_checkpoint_path is not None: self.predictor = build_sam2_video_predictor(model_cfg, sam_checkpoint_path, device="cuda") if flux_kontext_path is not None: self.flux_kontext = FluxKontextPipeline.from_pretrained(flux_kontext_path, torch_dtype=torch.bfloat16).to("cuda") def __call__(self, video_path, refer_image_path, output_path, resolution_area=[1280, 720], fps=30, iterations=3, k=7, w_len=1, h_len=1, retarget_flag=False, use_flux=False, replace_flag=False, pts_by_frame=None, lbs_by_frame=None): if replace_flag: frames = get_frames(video_path, resolution_area, fps) height, width = frames[0].shape[:2] if not pts_by_frame and not lbs_by_frame: ############################################################################ tpl_pose_metas = self.pose2d(frames) face_images = [] for idx, meta in enumerate(tpl_pose_metas): face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3, image_shape=(frames[0].shape[0], frames[0].shape[1])) x1, x2, y1, y2 = face_bbox_for_image face_image = frames[idx][y1:y2, x1:x2] face_image = cv2.resize(face_image, (512, 512)) face_images.append(face_image) logger.info(f"Processing reference image: {refer_image_path}") refer_img = cv2.imread(refer_image_path) src_ref_path = os.path.join(output_path, 'src_ref.png') shutil.copy(refer_image_path, src_ref_path) refer_img = refer_img[..., ::-1] refer_img = padding_resize(refer_img, height, width) logger.info(f"Processing template video: {video_path}") tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas] cond_images = [] for idx, meta in enumerate(tpl_retarget_pose_metas): canvas = np.zeros_like(refer_img) conditioning_image = draw_aapose_by_meta_new(canvas, meta) cond_images.append(conditioning_image) ############################################################################ ############################################################################ masks = self.get_mask_from_face_bbox(frames, 400, tpl_pose_metas) bg_images = [] aug_masks = [] for frame, mask in zip(frames, masks): if iterations > 0: _, each_mask = get_mask_body_img(frame, mask, iterations=iterations, k=k) each_aug_mask = get_aug_mask(each_mask, w_len=w_len, h_len=h_len) else: each_aug_mask = mask each_bg_image = frame * (1 - each_aug_mask[:, :, None]) bg_images.append(each_bg_image) aug_masks.append(each_aug_mask) ############################################################################ else: ############################################################################ masks = self.get_mask_from_face_bbox_v2(frames, pts_by_frame=pts_by_frame, lbs_by_frame=lbs_by_frame) bg_images = [] aug_masks = [] for frame, mask in zip(frames, masks): if iterations > 0: _, each_mask = get_mask_body_img(frame, mask, iterations=iterations, k=k) # each_aug_mask = get_aug_mask(each_mask, w_len=w_len, h_len=h_len) each_aug_mask = quantize_mask_blocky(each_mask, block_w=16, block_h=16, occupancy=0.15) # each_aug_mask = each_mask else: each_aug_mask = mask each_bg_image = frame * (1 - each_aug_mask[:, :, None]) bg_images.append(each_bg_image) aug_masks.append(each_aug_mask) ############################################################################ ############################################################################ tpl_pose_metas = self.pose2d( frames, bbx=masks, # your per-frame masks list/array ) face_images = [] for idx, meta in enumerate(tpl_pose_metas): face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3, image_shape=(frames[0].shape[0], frames[0].shape[1])) x1, x2, y1, y2 = face_bbox_for_image face_image = frames[idx][y1:y2, x1:x2] face_image = cv2.resize(face_image, (512, 512)) face_images.append(face_image) logger.info(f"Processing reference image: {refer_image_path}") refer_img = cv2.imread(refer_image_path) src_ref_path = os.path.join(output_path, 'src_ref.png') shutil.copy(refer_image_path, src_ref_path) refer_img = refer_img[..., ::-1] refer_img = padding_resize(refer_img, height, width) logger.info(f"Processing template video: {video_path}") tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas] cond_images = [] for idx, meta in enumerate(tpl_retarget_pose_metas): canvas = np.zeros_like(refer_img) conditioning_image = draw_aapose_by_meta_new(canvas, meta) cond_images.append(conditioning_image) ############################################################################ src_face_path = os.path.join(output_path, 'src_face.mp4') mpy.ImageSequenceClip(face_images, fps=fps).write_videofile(src_face_path, logger=None) src_pose_path = os.path.join(output_path, 'src_pose.mp4') mpy.ImageSequenceClip(cond_images, fps=fps).write_videofile(src_pose_path, logger=None) src_bg_path = os.path.join(output_path, 'src_bg.mp4') mpy.ImageSequenceClip(bg_images, fps=fps).write_videofile(src_bg_path, logger=None) aug_masks_new = [np.stack([mask * 255, mask * 255, mask * 255], axis=2) for mask in aug_masks] src_mask_path = os.path.join(output_path, 'src_mask.mp4') mpy.ImageSequenceClip(aug_masks_new, fps=fps).write_videofile(src_mask_path, logger=None) return True else: logger.info(f"Processing reference image: {refer_image_path}") refer_img = cv2.imread(refer_image_path) src_ref_path = os.path.join(output_path, 'src_ref.png') shutil.copy(refer_image_path, src_ref_path) refer_img = refer_img[..., ::-1] refer_img = resize_by_area(refer_img, resolution_area[0] * resolution_area[1], divisor=16) refer_pose_meta = self.pose2d([refer_img])[0] logger.info(f"Processing template video: {video_path}") video_reader = VideoReader(video_path) frame_num = len(video_reader) video_fps = video_reader.get_avg_fps() # TODO: Maybe we can switch to PyAV later, which can get accurate frame num duration = video_reader.get_frame_timestamp(-1)[-1] expected_frame_num = int(duration * video_fps + 0.5) ratio = abs((frame_num - expected_frame_num)/frame_num) if ratio > 0.1: print("Warning: The difference between the actual number of frames and the expected number of frames is two large") frame_num = expected_frame_num if fps == -1: fps = video_fps target_num = int(frame_num / video_fps * fps) idxs = get_frame_indices(frame_num, video_fps, target_num, fps) frames = video_reader.get_batch(idxs).asnumpy() logger.info(f"Processing pose meta") tpl_pose_meta0 = self.pose2d(frames[:1])[0] tpl_pose_metas = self.pose2d(frames) face_images = [] for idx, meta in enumerate(tpl_pose_metas): face_bbox_for_image = get_face_bboxes(meta['keypoints_face'][:, :2], scale=1.3, image_shape=(frames[0].shape[0], frames[0].shape[1])) x1, x2, y1, y2 = face_bbox_for_image face_image = frames[idx][y1:y2, x1:x2] face_image = cv2.resize(face_image, (512, 512)) face_images.append(face_image) if retarget_flag: if use_flux: tpl_prompt, refer_prompt = self.get_editing_prompts(tpl_pose_metas, refer_pose_meta) refer_input = Image.fromarray(refer_img) refer_edit = self.flux_kontext( image=refer_input, height=refer_img.shape[0], width=refer_img.shape[1], prompt=refer_prompt, guidance_scale=2.5, num_inference_steps=28, ).images[0] refer_edit = Image.fromarray(padding_resize(np.array(refer_edit), refer_img.shape[0], refer_img.shape[1])) refer_edit_path = os.path.join(output_path, 'refer_edit.png') refer_edit.save(refer_edit_path) refer_edit_pose_meta = self.pose2d([np.array(refer_edit)])[0] tpl_img = frames[1] tpl_input = Image.fromarray(tpl_img) tpl_edit = self.flux_kontext( image=tpl_input, height=tpl_img.shape[0], width=tpl_img.shape[1], prompt=tpl_prompt, guidance_scale=2.5, num_inference_steps=28, ).images[0] tpl_edit = Image.fromarray(padding_resize(np.array(tpl_edit), tpl_img.shape[0], tpl_img.shape[1])) tpl_edit_path = os.path.join(output_path, 'tpl_edit.png') tpl_edit.save(tpl_edit_path) tpl_edit_pose_meta0 = self.pose2d([np.array(tpl_edit)])[0] tpl_retarget_pose_metas = get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, tpl_edit_pose_meta0, refer_edit_pose_meta) else: tpl_retarget_pose_metas = get_retarget_pose(tpl_pose_meta0, refer_pose_meta, tpl_pose_metas, None, None) else: tpl_retarget_pose_metas = [AAPoseMeta.from_humanapi_meta(meta) for meta in tpl_pose_metas] cond_images = [] for idx, meta in enumerate(tpl_retarget_pose_metas): if retarget_flag: canvas = np.zeros_like(refer_img) conditioning_image = draw_aapose_by_meta_new(canvas, meta) else: canvas = np.zeros_like(frames[0]) conditioning_image = draw_aapose_by_meta_new(canvas, meta) conditioning_image = padding_resize(conditioning_image, refer_img.shape[0], refer_img.shape[1]) cond_images.append(conditioning_image) src_face_path = os.path.join(output_path, 'src_face.mp4') mpy.ImageSequenceClip(face_images, fps=fps).write_videofile(src_face_path, logger=None) src_pose_path = os.path.join(output_path, 'src_pose.mp4') mpy.ImageSequenceClip(cond_images, fps=fps).write_videofile(src_pose_path, logger=None) return True def get_editing_prompts(self, tpl_pose_metas, refer_pose_meta): arm_visible = False leg_visible = False for tpl_pose_meta in tpl_pose_metas: tpl_keypoints = tpl_pose_meta['keypoints_body'] if tpl_keypoints[3].all() != 0 or tpl_keypoints[4].all() != 0 or tpl_keypoints[6].all() != 0 or tpl_keypoints[7].all() != 0: if (tpl_keypoints[3][0] <= 1 and tpl_keypoints[3][1] <= 1 and tpl_keypoints[3][2] >= 0.75) or (tpl_keypoints[4][0] <= 1 and tpl_keypoints[4][1] <= 1 and tpl_keypoints[4][2] >= 0.75) or \ (tpl_keypoints[6][0] <= 1 and tpl_keypoints[6][1] <= 1 and tpl_keypoints[6][2] >= 0.75) or (tpl_keypoints[7][0] <= 1 and tpl_keypoints[7][1] <= 1 and tpl_keypoints[7][2] >= 0.75): arm_visible = True if tpl_keypoints[9].all() != 0 or tpl_keypoints[12].all() != 0 or tpl_keypoints[10].all() != 0 or tpl_keypoints[13].all() != 0: if (tpl_keypoints[9][0] <= 1 and tpl_keypoints[9][1] <= 1 and tpl_keypoints[9][2] >= 0.75) or (tpl_keypoints[12][0] <= 1 and tpl_keypoints[12][1] <= 1 and tpl_keypoints[12][2] >= 0.75) or \ (tpl_keypoints[10][0] <= 1 and tpl_keypoints[10][1] <= 1 and tpl_keypoints[10][2] >= 0.75) or (tpl_keypoints[13][0] <= 1 and tpl_keypoints[13][1] <= 1 and tpl_keypoints[13][2] >= 0.75): leg_visible = True if arm_visible and leg_visible: break if leg_visible: if tpl_pose_meta['width'] > tpl_pose_meta['height']: tpl_prompt = "Change the person to a standard T-pose (facing forward with arms extended). The person is standing. Feet and Hands are visible in the image." else: tpl_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. The person is standing. Feet and Hands are visible in the image." if refer_pose_meta['width'] > refer_pose_meta['height']: refer_prompt = "Change the person to a standard T-pose (facing forward with arms extended). The person is standing. Feet and Hands are visible in the image." else: refer_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. The person is standing. Feet and Hands are visible in the image." elif arm_visible: if tpl_pose_meta['width'] > tpl_pose_meta['height']: tpl_prompt = "Change the person to a standard T-pose (facing forward with arms extended). Hands are visible in the image." else: tpl_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. Hands are visible in the image." if refer_pose_meta['width'] > refer_pose_meta['height']: refer_prompt = "Change the person to a standard T-pose (facing forward with arms extended). Hands are visible in the image." else: refer_prompt = "Change the person to a standard pose with the face oriented forward and arms extending straight down by the sides. Hands are visible in the image." else: tpl_prompt = "Change the person to face forward." refer_prompt = "Change the person to face forward." return tpl_prompt, refer_prompt def get_mask_from_face_bbox_v2( self, frames, pts_by_frame: dict[int, list[list[float]]] | None = None, lbs_by_frame: dict[int, list[int | float]] | None = None, ): """ Args: frames: list/array of HxWx3 uint8 frames. pts_by_frame: {frame_idx: [[x,y], ...], ...} labels_by_frame: {frame_idx: [0/1,...], ...} Returns: all_mask: list[np.uint8 mask] of length len(frames), each (H, W) in {0,1} """ print(f"lbs_by_frame:{lbs_by_frame}") print(f"pts_by_frame:{pts_by_frame}") # --- safety & normalization --- if pts_by_frame is None: pts_by_frame = {} if lbs_by_frame is None: lbs_by_frame = {} # normalize keys to int (in case they arrived as strings) pts_by_frame = {int(k): v for k, v in pts_by_frame.items()} lbs_by_frame = {int(k): v for k, v in lbs_by_frame.items()} H, W = frames[0].shape[:2] device = "cuda" if torch.cuda.is_available() else "cpu" with torch.autocast(device_type=device, dtype=torch.bfloat16 if device == "cuda" else torch.float16): # 1) init SAM2 video predictor state inference_state = self.predictor.init_state(images=np.array(frames), device=device) # 2) feed all per-frame clicks before propagating # We use the *same obj_id* (0) so all clicks describe one object, # no matter which frame they were added on. for fidx in sorted(pts_by_frame.keys()): pts = np.array(pts_by_frame.get(fidx, []), dtype=np.float32) lbs = np.array(lbs_by_frame.get(fidx, []), dtype=np.int32) if pts.size == 0: continue # nothing to add for this frame # (optional) sanity: make sure lens match if len(pts) != len(lbs): raise ValueError(f"Points/labels length mismatch at frame {fidx}: {len(pts)} vs {len(lbs)}") self.predictor.add_new_points( inference_state=inference_state, frame_idx=int(fidx), obj_id=0, points=pts, labels=lbs, ) # 3) propagate across the whole video video_segments = {} for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video( inference_state, start_frame_idx=0 ): # store boolean masks per object id for this frame video_segments[out_frame_idx] = { out_obj_id: (out_mask_logits[i] > 0.0).to("cpu").numpy() for i, out_obj_id in enumerate(out_obj_ids) } # 4) collect masks in order; fall back to zeros where predictor returned nothing all_mask = [] zero_mask = np.zeros((H, W), dtype=np.uint8) for out_frame_idx in range(len(frames)): if out_frame_idx in video_segments and len(video_segments[out_frame_idx]) > 0: mask = next(iter(video_segments[out_frame_idx].values())) if mask.ndim == 3: # (1, H, W) -> (H, W) mask = mask[0] mask = mask.astype(np.uint8) else: mask = zero_mask all_mask.append(mask) return all_mask def get_mask_from_face_bbox(self, frames, th_step, kp2ds_all): """ Build masks using a face bounding box per key frame (derived from keypoints_face), then propagate with SAM2 across each chunk of frames. """ H, W = frames[0].shape[:2] def _clip_box(x1, y1, x2, y2, W, H): x1 = max(0, min(int(x1), W - 1)) x2 = max(0, min(int(x2), W - 1)) y1 = max(0, min(int(y1), H - 1)) y2 = max(0, min(int(y2), H - 1)) if x2 <= x1: x2 = min(W - 1, x1 + 1) if y2 <= y1: y2 = min(H - 1, y1 + 1) return x1, y1, x2, y2 frame_num = len(frames) if frame_num < th_step: num_step = 1 else: num_step = (frame_num + th_step) // th_step all_mask = [] for step_idx in range(num_step): each_frames = frames[step_idx * th_step:(step_idx + 1) * th_step] kp2ds = kp2ds_all[step_idx * th_step:(step_idx + 1) * th_step] if len(each_frames) == 0: continue # pick a few key frames in this chunk key_frame_num = 4 if len(each_frames) > 4 else 1 key_frame_step = max(1, len(kp2ds) // key_frame_num) key_frame_index_list = list(range(0, len(kp2ds), key_frame_step))[:key_frame_num] # compute face boxes on the selected key frames key_frame_boxes = [] for kfi in key_frame_index_list: meta = kp2ds[kfi] # get_face_bboxes returns (x1, x2, y1, y2) in your code x1, x2, y1, y2 = get_face_bboxes( meta['keypoints_face'][:, :2], scale=1.3, image_shape=(H, W) ) x1, y1, x2, y2 = _clip_box(x1, y1, x2, y2, W, H) key_frame_boxes.append(np.array([x1, y1, x2, y2], dtype=np.float32)) # init SAM2 for this chunk with torch.autocast(device_type="cuda", dtype=torch.bfloat16): inference_state = self.predictor.init_state(images=np.array(each_frames), device="cuda") self.predictor.reset_state(inference_state) ann_obj_id = 1 # seed with box prompts (preferred), else fall back to points for ann_frame_idx, box_xyxy in zip(key_frame_index_list, key_frame_boxes): used_box = False try: # If your predictor exposes a box API, this is ideal. _ = self.predictor.add_new_box( inference_state=inference_state, frame_idx=ann_frame_idx, obj_id=ann_obj_id, box=box_xyxy[None, :] # shape (1, 4) ) used_box = True except Exception: used_box = False if not used_box: # Fallback: sample a few positive points inside the box x1, y1, x2, y2 = box_xyxy.astype(int) cx, cy = (x1 + x2) // 2, (y1 + y2) // 2 pts = np.array([ [cx, cy], [x1 + (x2 - x1) // 4, cy], [x2 - (x2 - x1) // 4, cy], [cx, y1 + (y2 - y1) // 4], [cx, y2 - (y2 - y1) // 4], ], dtype=np.int32) labels = np.ones(len(pts), dtype=np.int32) # 1 = positive _ = self.predictor.add_new_points( inference_state=inference_state, frame_idx=ann_frame_idx, obj_id=ann_obj_id, points=pts, labels=labels, ) # propagate across the chunk video_segments = {} for out_frame_idx, out_obj_ids, out_mask_logits in self.predictor.propagate_in_video(inference_state): video_segments[out_frame_idx] = { out_obj_id: (out_mask_logits[i] > 0.0).cpu().numpy() for i, out_obj_id in enumerate(out_obj_ids) } # collect masks (single object id) for out_frame_idx in range(len(video_segments)): # (H, W) boolean/uint8 mask = next(iter(video_segments[out_frame_idx].values())) mask = mask[0].astype(np.uint8) all_mask.append(mask) return all_mask