revert: remove transparent option, keep JPG/MP4 for performance
- bg_color options: green/blue/black/white (default: green) - Previews saved as JPG, SelectActorPreview outputs [N,H,W,3] RGB - Removed merge_segments_rgba() and all RGBA/WebM logic - 3-channel BGR throughout — no alpha overhead
This commit is contained in:
Cherokee
+29
-53
@@ -49,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, merge_segments_rgba, generate_actor_json
|
||||
from pipeline.merger import merge_segments, generate_actor_json
|
||||
|
||||
# PersonSegmenter may not exist yet (parallel task creating it)
|
||||
try:
|
||||
@@ -310,8 +310,8 @@ class VideoActorExtractor:
|
||||
},
|
||||
),
|
||||
"bg_color": (
|
||||
["transparent", "green", "blue", "black", "white"],
|
||||
{"default": "transparent"},
|
||||
["green", "blue", "black", "white"],
|
||||
{"default": "green"},
|
||||
),
|
||||
},
|
||||
}
|
||||
@@ -331,7 +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",
|
||||
bg_color: str = "green",
|
||||
) -> Tuple[str, str, int]:
|
||||
"""
|
||||
Run the full actor extraction pipeline.
|
||||
@@ -753,17 +753,17 @@ class VideoActorExtractor:
|
||||
top5 = sorted(
|
||||
face_detected_frames, key=lambda x: actor_face_areas[x[0]], reverse=True
|
||||
)[:5]
|
||||
# Previews always saved as PNG with alpha channel
|
||||
# Previews saved as JPG with selected bg color
|
||||
bg_bgr = BG_COLOR_MAP.get(bg_color, BG_COLOR_GREEN)
|
||||
preview_frame_indexes = []
|
||||
for k, (fi, bool_mask, _, _) in enumerate(top5):
|
||||
frame_bgr = frame_lookup.get(fi)
|
||||
if frame_bgr is None:
|
||||
continue
|
||||
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)
|
||||
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
|
||||
@@ -784,50 +784,26 @@ class VideoActorExtractor:
|
||||
if not segments_data:
|
||||
continue
|
||||
|
||||
# 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)
|
||||
# Convert bool-mask segments to BGR 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:
|
||||
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}.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,
|
||||
)
|
||||
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)
|
||||
|
||||
+9
-25
@@ -42,46 +42,30 @@ class SelectActorPreview:
|
||||
|
||||
Returns:
|
||||
Tuple of:
|
||||
- 4D tensor [N, H, W, 4] RGBA float32 in 0-1 range.
|
||||
- 4D tensor [N, H, W, 3] RGB float32 in 0-1 range.
|
||||
- List of frame indexes: [frame_idx_0, frame_idx_1, ...]
|
||||
"""
|
||||
import cv2
|
||||
|
||||
# Previews are now always PNG with alpha channel
|
||||
pattern = os.path.join(output_dir, "previews", f"actor_{actor_index}_*.png")
|
||||
pattern = os.path.join(output_dir, "previews", f"actor_{actor_index}_*.jpg")
|
||||
paths = sorted(glob.glob(pattern))
|
||||
|
||||
if not paths:
|
||||
# Fallback: check for legacy .jpg previews
|
||||
pattern_jpg = os.path.join(
|
||||
output_dir, "previews", f"actor_{actor_index}_*.jpg"
|
||||
)
|
||||
paths = sorted(glob.glob(pattern_jpg))
|
||||
|
||||
if not paths:
|
||||
print(f"[SelectActorPreview] No previews found: {pattern}")
|
||||
return (torch.zeros(1, 512, 512, 4, dtype=torch.float32), [])
|
||||
return (torch.zeros(1, 512, 512, 3, dtype=torch.float32), [])
|
||||
|
||||
frames = []
|
||||
for p in paths:
|
||||
img_bgra = cv2.imread(p, cv2.IMREAD_UNCHANGED)
|
||||
if img_bgra is None:
|
||||
img_bgr = cv2.imread(p)
|
||||
if img_bgr is None:
|
||||
continue
|
||||
if img_bgra.ndim == 2:
|
||||
# Grayscale — convert to BGRA
|
||||
img_bgra = cv2.cvtColor(img_bgra, cv2.COLOR_GRAY2BGRA)
|
||||
elif img_bgra.shape[2] == 3:
|
||||
# BGR without alpha — add full opacity
|
||||
alpha = np.full((*img_bgra.shape[:2], 255), dtype=np.uint8)
|
||||
img_bgra = np.dstack([img_bgra, alpha])
|
||||
# BGRA -> RGBA
|
||||
img_rgba = img_bgra[:, :, [2, 1, 0, 3]].copy()
|
||||
frames.append(img_rgba.astype(np.float32) / 255.0)
|
||||
img_rgb = img_bgr[:, :, ::-1].copy()
|
||||
frames.append(img_rgb.astype(np.float32) / 255.0)
|
||||
|
||||
if not frames:
|
||||
return (torch.zeros(1, 512, 512, 4, dtype=torch.float32), [])
|
||||
return (torch.zeros(1, 512, 512, 3, dtype=torch.float32), [])
|
||||
|
||||
tensor = torch.from_numpy(np.stack(frames)) # [N, H, W, 4]
|
||||
tensor = torch.from_numpy(np.stack(frames)) # [N, H, W, 3]
|
||||
|
||||
# Read frame indexes from the JSON file written by VideoActorExtractor
|
||||
indexes_path = os.path.join(
|
||||
|
||||
@@ -152,75 +152,6 @@ def merge_segments(
|
||||
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,
|
||||
|
||||
Reference in New Issue
Block a user