File size: 2,926 Bytes
fb1ae50 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 | """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) # B,H,W bool
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
# fit aspect
if bw / bh > target_ar:
bh = bw / target_ar
else:
bw = bh * target_ar
# clamp size to frame
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) # B,h,w,1
region = out[:, y0:y1, x0:x1, :]
out[:, y0:y1, x0:x1, :] = gen * alpha_t + region * (1 - alpha_t)
return out
|