nodes/ahkimkoo-ComfyUI-VideoActorExtract
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add bg_color option for cutout background (transparent/green/blue/black/white)

Browse Source 
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.
This commit is contained in:
Cherokee
2026-04-30 16:35:17 +08:00
parent 2f58b7e589
commit 85e2dc2fd2
3 changed files with 213 additions and 74 deletions
Download Patch File Download Diff File Expand all files Collapse all files
core/config.py
+13
View File
@@ -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
nodes/actor_extractor.py
+64 -28
View File
@@ -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)
pipeline/merger.py
+136 -46
View File
@@ -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