From 85e2dc2fd25d19e476204ceccb99ef2aade4f3f5 Mon Sep 17 00:00:00 2001 From: Cherokee Date: Thu, 30 Apr 2026 16:35:17 +0800 Subject: [PATCH] feat: add bg_color option for cutout background (transparent/green/blue/black/white) Default is transparent (WebM with VP9 alpha). Colored backgrounds output MP4 with the selected solid color filling non-person pixels. Transparent previews saved as PNG, colored as JPG. --- core/config.py | 13 +++ nodes/actor_extractor.py | 92 ++++++++++++++------ pipeline/merger.py | 182 +++++++++++++++++++++++++++++---------- 3 files changed, 213 insertions(+), 74 deletions(-) diff --git a/core/config.py b/core/config.py index cd6a8de..72baca0 100644 --- a/core/config.py +++ b/core/config.py @@ -27,6 +27,19 @@ YOLO_INFER_SIZE = 640 # Green screen GREEN_SCREEN_COLOR = (0, 255, 0) # BGR pure green +# Background colors (BGR format) +BG_COLOR_GREEN = (0, 255, 0) +BG_COLOR_BLUE = (255, 0, 0) +BG_COLOR_BLACK = (0, 0, 0) +BG_COLOR_WHITE = (255, 255, 255) + +BG_COLOR_MAP = { + "green": BG_COLOR_GREEN, + "blue": BG_COLOR_BLUE, + "black": BG_COLOR_BLACK, + "white": BG_COLOR_WHITE, +} + # Segment gap (green screen frames between merged segments) SEGMENT_GAP_SEC = 0.5 SEGMENT_GAP_FRAMES_DEFAULT = 15 # will be computed from fps diff --git a/nodes/actor_extractor.py b/nodes/actor_extractor.py index f8e687a..1dfb744 100644 --- a/nodes/actor_extractor.py +++ b/nodes/actor_extractor.py @@ -36,6 +36,8 @@ from core.config import ( DEFAULT_YOLO_MODEL, DEFAULT_MAX_LOST_FRAMES, SEGMENT_GAP_SEC, + BG_COLOR_MAP, + BG_COLOR_GREEN, ) # DEFAULT_SEG_MODEL may not exist yet (parallel task adding it) @@ -47,7 +49,7 @@ from pipeline.identity import IdentityCluster # face_threshold is kept for backward compatibility but identity clustering # now uses spatial-temporal constraints to prevent merging co-occurring tracks. -from pipeline.merger import merge_segments, generate_actor_json +from pipeline.merger import merge_segments, merge_segments_rgba, generate_actor_json # PersonSegmenter may not exist yet (parallel task creating it) try: @@ -307,6 +309,10 @@ class VideoActorExtractor: "tooltip": "Process every Nth frame. 1=all frames, 2=every other frame, etc. Higher values are faster but less precise.", }, ), + "bg_color": ( + ["transparent", "green", "blue", "black", "white"], + {"default": "transparent"}, + ), }, } @@ -325,6 +331,7 @@ class VideoActorExtractor: face_threshold: float = DEFAULT_FACE_THRESHOLD, min_track_length: int = DEFAULT_MIN_TRACK_LENGTH, skip_every_n: int = 1, + bg_color: str = "transparent", ) -> Tuple[str, str, int]: """ Run the full actor extraction pipeline. @@ -751,10 +758,19 @@ class VideoActorExtractor: frame_bgr = frame_lookup.get(fi) if frame_bgr is None: continue - masked_bgr = frame_bgr.copy() - masked_bgr[~bool_mask] = (0, 255, 0) - preview_path = os.path.join(preview_dir, f"actor_{i}_{k}.jpg") - cv2.imwrite(preview_path, masked_bgr) + if bg_color == "transparent": + # Save as PNG with alpha channel + bgra = np.zeros((img_h, img_w, 4), dtype=np.uint8) + bgra[bool_mask, :3] = frame_bgr[bool_mask] + bgra[bool_mask, 3] = 255 + preview_path = os.path.join(preview_dir, f"actor_{i}_{k}.png") + cv2.imwrite(preview_path, bgra) + else: + bg_bgr = BG_COLOR_MAP.get(bg_color, BG_COLOR_GREEN) + masked_bgr = frame_bgr.copy() + masked_bgr[~bool_mask] = bg_bgr + preview_path = os.path.join(preview_dir, f"actor_{i}_{k}.jpg") + cv2.imwrite(preview_path, masked_bgr) preview_frame_indexes.append(fi) # Save frame indexes for SelectActorPreview to read @@ -775,30 +791,50 @@ class VideoActorExtractor: if not segments_data: continue - # Convert bool-mask segments to BGR green-screened frames for encoding - segments_bgr: List[List[np.ndarray]] = [] - for seg in segments_data: - bgr_frames = [] - for fi, bool_mask, _ in seg: - frame_bgr = frame_lookup.get(fi) - if frame_bgr is None: - # Fallback: pure green frame - frame_bgr = np.full( - (img_h, img_w, 3), (0, 255, 0), dtype=np.uint8 - ) - greened = frame_bgr.copy() - greened[~bool_mask] = (0, 255, 0) - bgr_frames.append(greened) - segments_bgr.append(bgr_frames) + # Convert bool-mask segments to frames for encoding + if bg_color == "transparent": + # Create RGBA frames: person pixels opaque, background transparent + segments_rgba: List[List[np.ndarray]] = [] + for seg in segments_data: + rgba_frames = [] + for fi, bool_mask, _ in seg: + frame_bgr = frame_lookup.get(fi) + if frame_bgr is None: + frame_bgr = np.zeros((img_h, img_w, 3), dtype=np.uint8) + bgra = np.zeros((img_h, img_w, 4), dtype=np.uint8) + bgra[bool_mask, :3] = frame_bgr[bool_mask] + bgra[bool_mask, 3] = 255 + rgba_frames.append(bgra) + segments_rgba.append(rgba_frames) - # Encode to video via merge_segments - output_video_path = os.path.join(output_dir, f"{actor_id}.mp4") - success = merge_segments( - segments_bgr, - fps=fps, - output_path=output_video_path, - gap_sec=SEGMENT_GAP_SEC, - ) + output_video_path = os.path.join(output_dir, f"{actor_id}.webm") + success = merge_segments_rgba( + segments_rgba, fps=fps, output_path=output_video_path + ) + else: + # Colored background compositing + bg_bgr = BG_COLOR_MAP.get(bg_color, BG_COLOR_GREEN) + segments_bgr: List[List[np.ndarray]] = [] + for seg in segments_data: + bgr_frames = [] + for fi, bool_mask, _ in seg: + frame_bgr = frame_lookup.get(fi) + if frame_bgr is None: + frame_bgr = np.full( + (img_h, img_w, 3), bg_bgr, dtype=np.uint8 + ) + composited = frame_bgr.copy() + composited[~bool_mask] = bg_bgr + bgr_frames.append(composited) + segments_bgr.append(bgr_frames) + + output_video_path = os.path.join(output_dir, f"{actor_id}.mp4") + success = merge_segments( + segments_bgr, + fps=fps, + output_path=output_video_path, + gap_sec=SEGMENT_GAP_SEC, + ) if success: total_output_frames = sum(s["frame_count"] for s in segments_info) diff --git a/pipeline/merger.py b/pipeline/merger.py index 794f137..c14675f 100644 --- a/pipeline/merger.py +++ b/pipeline/merger.py @@ -24,31 +24,34 @@ from core.config import ( def _get_ffmpeg_codec() -> str: """Select appropriate ffmpeg video codec for the current platform.""" import platform + system = platform.system() if system == "Darwin": return DEFAULT_VIDEO_CODEC_MAC return DEFAULT_VIDEO_CODEC -def _encode_frames_to_temp(frames: List[np.ndarray], temp_path: str, fps: float) -> bool: +def _encode_frames_to_temp( + frames: List[np.ndarray], temp_path: str, fps: float +) -> bool: """ Encode a list of frames to a temporary video file. - + Returns True on success, False on failure. """ if not frames: return False - + h, w = frames[0].shape[:2] codec = _get_ffmpeg_codec() - + try: fourcc = cv2.VideoWriter_fourcc(*"mp4v") out = cv2.VideoWriter(temp_path, fourcc, fps, (w, h)) if not out.isOpened(): # Fallback to ffmpeg directly return _encode_with_ffmpeg(frames, temp_path, fps) - + for frame in frames: out.write(frame) out.release() @@ -61,28 +64,41 @@ def _encode_with_ffmpeg(frames: List[np.ndarray], temp_path: str, fps: float) -> """Encode frames using ffmpeg pipe.""" if not frames: return False - + h, w = frames[0].shape[:2] codec = _get_ffmpeg_codec() - + cmd = [ - "ffmpeg", "-y", - "-f", "rawvideo", - "-vcodec", "rawvideo", - "-s", f"{w}x{h}", - "-pix_fmt", "bgr24", - "-r", str(fps), - "-i", "-", - "-c:v", codec, - "-pix_fmt", DEFAULT_PIXEL_FORMAT, - "-crf", str(DEFAULT_CRF), - "-preset", "fast", + "ffmpeg", + "-y", + "-f", + "rawvideo", + "-vcodec", + "rawvideo", + "-s", + f"{w}x{h}", + "-pix_fmt", + "bgr24", + "-r", + str(fps), + "-i", + "-", + "-c:v", + codec, + "-pix_fmt", + DEFAULT_PIXEL_FORMAT, + "-crf", + str(DEFAULT_CRF), + "-preset", + "fast", "-an", temp_path, ] - + try: - proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + proc = subprocess.Popen( + cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE + ) raw_bytes = b"".join([cv2.imencode(".png", f)[1].tobytes() for f in frames]) # Actually pipe raw BGR raw_data = np.stack(frames).tobytes() @@ -107,34 +123,104 @@ def merge_segments( ) -> bool: """ Merge multiple segments into one video. - + Args: segments: List of frame lists, each list is one continuous segment fps: Frame rate output_path: Output MP4 path gap_sec: (kept for API compatibility, no longer inserts green gaps) - + Returns: True on success """ if not segments or all(len(s) == 0 for s in segments): return False - + # Filter out empty segments segments = [s for s in segments if len(s) > 0] if not segments: return False - + if len(segments) == 1: return _encode_frames_to_temp(segments[0], output_path, fps) - + # Concatenate all segments without gaps merged = [] for seg in segments: merged.extend(seg) - + return _encode_frames_to_temp(merged, output_path, fps) + +def merge_segments_rgba( + segments: List[List[np.ndarray]], + fps: float, + output_path: str, +) -> bool: + """ + Merge RGBA frame segments into a WebM video with VP9 alpha channel. + + Each frame must be a BGRA uint8 numpy array (OpenCV's default RGBA + ordering). Person pixels have alpha=255, background pixels alpha=0. + + Args: + segments: List of frame lists, each list is one continuous segment. + Frames are BGRA uint8 arrays. + fps: Frame rate. + output_path: Output .webm file path. + + Returns: + True on success, False on failure. + """ + if not segments or all(len(s) == 0 for s in segments): + return False + + # Filter out empty segments and flatten + frames = [f for seg in segments if seg for f in seg] + if not frames: + return False + + h, w = frames[0].shape[:2] + + cmd = [ + "ffmpeg", + "-y", + "-f", + "rawvideo", + "-vcodec", + "rawvideo", + "-s", + f"{w}x{h}", + "-pix_fmt", + "bgra", + "-r", + str(fps), + "-i", + "-", + "-c:v", + "libvpx-vp9", + "-pix_fmt", + "yuva420p", + "-auto-alt-ref", + "0", + "-an", + output_path, + ] + + try: + proc = subprocess.Popen( + cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + ) + raw_data = np.stack(frames).tobytes() + proc.communicate(input=raw_data, timeout=300) + return proc.returncode == 0 + except Exception: + return False + + def generate_actor_json( actor_data: Dict[str, dict], video_info: dict, @@ -142,7 +228,7 @@ def generate_actor_json( ) -> str: """ Generate actor_info.json file. - + Args: actor_data: { "actor_0": { @@ -157,7 +243,7 @@ def generate_actor_json( } video_info: {"total_frames": ..., "fps": ..., "width": ..., "height": ..., "duration_sec": ...} output_path: Path to write JSON file - + Returns: Path to the generated JSON file """ @@ -165,30 +251,34 @@ def generate_actor_json( "actors": [], "video_info": video_info, } - + for actor_id in sorted(actor_data.keys()): info = actor_data[actor_id] segments = [] for i, seg in enumerate(info["segments"]): start_time = seg["start_frame"] / video_info["fps"] end_time = seg["end_frame"] / video_info["fps"] - segments.append({ - "segment_id": i, - "start_frame": seg["start_frame"], - "end_frame": seg["end_frame"], - "start_time_sec": round(start_time, 2), - "end_time_sec": round(end_time, 2), - "frame_count": seg["frame_count"], - }) - - result["actors"].append({ - "actor_id": actor_id, - "segment_count": info["segment_count"], - "total_frames": info["total_frames"], - "segments": segments, - }) - + segments.append( + { + "segment_id": i, + "start_frame": seg["start_frame"], + "end_frame": seg["end_frame"], + "start_time_sec": round(start_time, 2), + "end_time_sec": round(end_time, 2), + "frame_count": seg["frame_count"], + } + ) + + result["actors"].append( + { + "actor_id": actor_id, + "segment_count": info["segment_count"], + "total_frames": info["total_frames"], + "segments": segments, + } + ) + with open(output_path, "w", encoding="utf-8") as f: json.dump(result, f, indent=2, ensure_ascii=False) - + return output_path