| """Crop-around-person and composite-back helpers for person replacement. |
| |
| Uses the SAM3 pose mask to find the person bbox, crops the driving video |
| around them (so the generation model's centered output aligns), and pastes |
| the generated person back into the original frames. |
| """ |
|
|
| import cv2 |
| import numpy as np |
| import torch |
|
|
|
|
| def person_bbox(mask, thresh=0.1): |
| """Union bbox of the person across all mask frames. |
| |
| mask: [B, H, W, 3] colored mask (person = any bright channel). |
| Returns (x0, y0, x1, y1) or None if nothing found. |
| """ |
| m = (mask[..., :3].max(dim=-1).values > thresh) |
| ys, xs = torch.nonzero(m.any(dim=0), as_tuple=True) |
| if len(xs) == 0: |
| return None |
| return int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1 |
|
|
|
|
| def fit_crop_box(bbox, frame_w, frame_h, target_ar, margin=0.15): |
| """Expand bbox by margin, adjust to target aspect ratio, clamp to frame.""" |
| x0, y0, x1, y1 = bbox |
| bw, bh = x1 - x0, y1 - y0 |
| cx, cy = (x0 + x1) / 2, (y0 + y1) / 2 |
| bw *= 1 + margin |
| bh *= 1 + margin |
| |
| if bw / bh > target_ar: |
| bh = bw / target_ar |
| else: |
| bw = bh * target_ar |
| |
| bw = min(bw, frame_w) |
| bh = min(bh, frame_h) |
| x0 = int(round(cx - bw / 2)) |
| y0 = int(round(cy - bh / 2)) |
| x0 = max(0, min(x0, frame_w - int(bw))) |
| y0 = max(0, min(y0, frame_h - int(bh))) |
| return x0, y0, x0 + int(bw), y0 + int(bh) |
|
|
|
|
| def crop_video(video, box): |
| """Crop ComfyUI IMAGE tensor [B,H,W,C] to (x0,y0,x1,y1).""" |
| x0, y0, x1, y1 = box |
| return video[:, y0:y1, x0:x1, :] |
|
|
|
|
| def resize_video(video, target_w, target_h): |
| """Resize ComfyUI IMAGE tensor [B,H,W,C] with lanczos via cv2.""" |
| arr = (video.cpu().numpy() * 255).astype(np.uint8) |
| out = np.stack([ |
| cv2.resize(f, (target_w, target_h), interpolation=cv2.INTER_LANCZOS4) |
| for f in arr |
| ]) |
| return torch.from_numpy(out).float() / 255.0 |
|
|
|
|
| def composite_back(original, generated, box, mask=None, feather=8): |
| """Paste generated person back into original frames. |
| |
| original: [B, H, W, 3] float 0-1 (source video) |
| generated: [B, h, w, 3] float 0-1 (generated, same AR as box) |
| box: (x0, y0, x1, y1) paste location |
| mask: optional [B, h, w] float 0-1 person alpha (None = paste all) |
| """ |
| x0, y0, x1, y1 = box |
| bw, bh = x1 - x0, y1 - y0 |
| gen = resize_video(generated, bw, bh) |
| out = original.clone() |
| if mask is None: |
| out[:, y0:y1, x0:x1, :] = gen |
| else: |
| alpha = mask.cpu().numpy() |
| if feather > 0: |
| alpha = np.stack([ |
| cv2.GaussianBlur(a, (feather * 2 + 1,) * 2, 0) for a in alpha |
| ]) |
| alpha_t = torch.from_numpy(alpha).float().unsqueeze(-1) |
| region = out[:, y0:y1, x0:x1, :] |
| out[:, y0:y1, x0:x1, :] = gen * alpha_t + region * (1 - alpha_t) |
| return out |
|
|
