Buckets:

fishxinyu
/

flexcombine-bench

Files

xet

fishxinyu/flexcombine-bench / tools /segment_human.py

fishxinyu

2 months ago

download

raw

9.22 kB

#!/usr/bin/env python3
"""
Extract person mask videos from videos using SAM3 text-prompted segmentation.
Tracks every detected person instance and saves a binary mask video per clip
(white = any person, black = background, H.264-encoded).  If multiple people
appear in a frame, their masks are merged into that frame.
Usage:
  # Process all .mp4 files found in video_dir:
  python segment_human.py \
    --video_dir /path/to/videos \
    --output_dir /path/to/output
  # Or restrict to clips listed in a CSV (needs a 'clip_id' column):
  python segment_human.py \
    --video_dir /path/to/videos \
    --csv /path/to/clips.csv \
    --output_dir /path/to/output
Output structure:
    output_dir/
      <clip_id>/
        masks.mp4   # binary mask video: white=persons, black=background, H.264
"""

import argparse
import csv
import glob
import os
import subprocess
import sys
from pathlib import Path

import cv2
import numpy as np


# ---------------------------------------------------------------------------
# Utilities
# ---------------------------------------------------------------------------

def get_video_fps(video_path: str) -> float:
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS)
    cap.release()
    return fps if fps > 0 else 30.0


def load_video_frames(video_path: str) -> list:
    """Return list of RGB numpy arrays (H, W, 3)."""
    cap = cv2.VideoCapture(video_path)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
    cap.release()
    return frames


def mask_to_numpy(mask) -> np.ndarray:
    """Convert a SAM3 binary mask (tensor or ndarray) to a bool (H, W) array."""
    if hasattr(mask, "numpy"):
        return mask.numpy().astype(bool)
    return np.asarray(mask, dtype=bool)


def save_mask_video(mask_frames: list, out_path: str, fps: float) -> None:
    """Write binary mask frames as an H.264 MP4.
    mask_frames: list of (H, W) bool arrays; True = person, False = background.
    Encodes via a temporary raw MP4 then re-encodes with ffmpeg libx264.
    """
    if not mask_frames:
        return
    h, w = mask_frames[0].shape[:2]
    tmp = out_path + ".tmp.mp4"
    writer = cv2.VideoWriter(tmp, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
    for mask in mask_frames:
        gray = (mask * 255).astype(np.uint8)
        bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
        writer.write(bgr)
    writer.release()
    ret = subprocess.run(
        ["ffmpeg", "-y", "-i", tmp, "-c:v", "libx264", "-pix_fmt", "yuv420p", out_path],
        capture_output=True,
    )
    os.remove(tmp)
    if ret.returncode != 0:
        print(f"  ffmpeg H.264 encode failed:\n{ret.stderr.decode()}")


# ---------------------------------------------------------------------------
# Core processing
# ---------------------------------------------------------------------------

def process_video(
    predictor,
    video_path: str,
    out_path: str,
    fps: float,
) -> None:
    frames = load_video_frames(video_path)
    if not frames:
        print(f"  Warning: no frames loaded from {video_path}")
        return

    # Start session
    response = predictor.handle_request(
        request=dict(type="start_session", resource_path=video_path)
    )
    session_id = response["session_id"]

    try:
        predictor.handle_request(
            request=dict(type="reset_session", session_id=session_id)
        )

        # Detect all person instances in frame 0
        response = predictor.handle_request(
            request=dict(
                type="add_prompt",
                session_id=session_id,
                frame_index=0,
                text="person",
            )
        )
        out0 = response["outputs"]
        obj_ids_0 = out0["out_obj_ids"].tolist()

        if not obj_ids_0:
            print(f"  No persons detected — skipping")
            return

        print(f"  Detected {len(obj_ids_0)} person(s): obj_ids={[int(x) for x in obj_ids_0]}")

        # Propagate through all frames — drain the stream before any numpy work
        # to avoid distributed-barrier deadlocks on multi-GPU setups.
        raw_outputs: dict[int, object] = {}
        for resp in predictor.handle_stream_request(
            request=dict(type="propagate_in_video", session_id=session_id)
        ):
            frame_idx = resp["frame_index"]
            if frame_idx < len(frames):
                raw_outputs[frame_idx] = resp["outputs"]

    finally:
        predictor.handle_request(
            request=dict(type="close_session", session_id=session_id)
        )

    # Build one combined mask per frame (OR across all detected persons).
    # Frames with no tracking output get a black (all-zero) mask.
    n = len(frames)
    mask_frames: list[np.ndarray] = []
    for frame_idx in range(n):
        h, w = frames[frame_idx].shape[:2]
        combined = np.zeros((h, w), dtype=bool)
        if frame_idx in raw_outputs:
            out = raw_outputs[frame_idx]
            for i in range(len(out["out_obj_ids"].tolist())):
                combined |= mask_to_numpy(out["out_binary_masks"][i])
        mask_frames.append(combined)

    save_mask_video(mask_frames, out_path, fps)
    print(f"  Saved mask video ({n} frames) → {out_path}")


# ---------------------------------------------------------------------------
# Entry point
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract person mask videos using SAM3. "
                    "Single-file mode: --input / --output. "
                    "Batch mode: --video_dir / --output_dir (and optionally --csv)."
    )
    # Single-file mode
    parser.add_argument("--input", default=None, help="Input .mp4 video file")
    parser.add_argument("--output", default=None, help="Output mask .mp4 file")
    # Batch mode
    parser.add_argument("--video_dir", default=None, help="Folder containing .mp4 files")
    parser.add_argument(
        "--csv", default=None,
        help="Optional CSV with a 'clip_id' column; if omitted, all .mp4s in "
             "video_dir are processed",
    )
    parser.add_argument(
        "--output_dir", default=None,
        help="Root directory for mask outputs (one masks.mp4 per clip subdirectory)",
    )
    parser.add_argument(
        "--checkpoint", default=None,
        help="Path to SAM 3 checkpoint .pt file (default: sam3/sam3/sam3.pt, "
             "or auto-download from HuggingFace)",
    )
    args = parser.parse_args()

    # Resolve mode and build (video_path, out_path) pairs
    entries: list[tuple[str, str]] = []  # (video_path, out_path)

    if args.input:
        if not args.output:
            parser.error("--input requires --output")
        entries = [(args.input, args.output)]
    elif args.video_dir:
        if not args.output_dir:
            parser.error("--video_dir requires --output_dir")
        if args.csv:
            with open(args.csv, newline="", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                for row in reader:
                    clip_id = row["clip_id"].strip()
                    video_path = os.path.join(args.video_dir, clip_id + ".mp4")
                    out_path = os.path.join(args.output_dir, clip_id, "masks.mp4")
                    entries.append((video_path, out_path))
            print(f"Loaded {len(entries)} entries from {args.csv}")
        else:
            mp4s = sorted(glob.glob(os.path.join(args.video_dir, "*.mp4")))
            for p in mp4s:
                clip_id = Path(p).stem
                out_path = os.path.join(args.output_dir, clip_id, "masks.mp4")
                entries.append((p, out_path))
            print(f"Found {len(entries)} .mp4 file(s) in {args.video_dir}")
    else:
        parser.error("Provide either --input/--output or --video_dir/--output_dir")

    if not entries:
        print("No videos to process — exiting.")
        sys.exit(0)

    # Load model
    print("Loading SAM 3 model …")
    import torch
    from sam3.model_builder import build_sam3_video_predictor

    checkpoint = args.checkpoint
    if checkpoint is None:
        script_dir = os.path.dirname(os.path.abspath(__file__))
        local = os.path.join(script_dir, "sam3", "sam3", "sam3.pt")
        checkpoint = local if os.path.exists(local) else None

    gpus_to_use = list(range(torch.cuda.device_count())) or [torch.cuda.current_device()]
    predictor = build_sam3_video_predictor(checkpoint_path=checkpoint, gpus_to_use=gpus_to_use)
    print("Model ready.\n")

    # Process
    for video_path, out_path in entries:
        if not os.path.exists(video_path):
            print(f"[SKIP] Not found: {video_path}")
            continue
        os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
        fps = get_video_fps(video_path)
        print(f"[{Path(video_path).stem}] fps={fps:.1f}  {video_path} → {out_path}")
        process_video(predictor=predictor, video_path=video_path, out_path=out_path, fps=fps)

    print("\nAll done.")


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 9.22 kB
Xet hash:: b8b6b227f7b19e047b9afbd0fe8a5eb26a29ec297ee730861d391ad29ed7c38c

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	#!/usr/bin/env python3
	"""
	Extract person mask videos from videos using SAM3 text-prompted segmentation.

	Tracks every detected person instance and saves a binary mask video per clip
	(white = any person, black = background, H.264-encoded). If multiple people
	appear in a frame, their masks are merged into that frame.

	Usage:
	# Process all .mp4 files found in video_dir:
	python segment_human.py \
	--video_dir /path/to/videos \
	--output_dir /path/to/output

	# Or restrict to clips listed in a CSV (needs a 'clip_id' column):
	python segment_human.py \
	--video_dir /path/to/videos \
	--csv /path/to/clips.csv \
	--output_dir /path/to/output

	Output structure:
	output_dir/
	<clip_id>/
	masks.mp4 # binary mask video: white=persons, black=background, H.264
	"""

	import argparse
	import csv
	import glob
	import os
	import subprocess
	import sys
	from pathlib import Path

	import cv2
	import numpy as np


	# ---------------------------------------------------------------------------
	# Utilities
	# ---------------------------------------------------------------------------

	def get_video_fps(video_path: str) -> float:
	cap = cv2.VideoCapture(video_path)
	fps = cap.get(cv2.CAP_PROP_FPS)
	cap.release()
	return fps if fps > 0 else 30.0


	def load_video_frames(video_path: str) -> list:
	"""Return list of RGB numpy arrays (H, W, 3)."""
	cap = cv2.VideoCapture(video_path)
	frames = []
	while True:
	ret, frame = cap.read()
	if not ret:
	break
	frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
	cap.release()
	return frames


	def mask_to_numpy(mask) -> np.ndarray:
	"""Convert a SAM3 binary mask (tensor or ndarray) to a bool (H, W) array."""
	if hasattr(mask, "numpy"):
	return mask.numpy().astype(bool)
	return np.asarray(mask, dtype=bool)


	def save_mask_video(mask_frames: list, out_path: str, fps: float) -> None:
	"""Write binary mask frames as an H.264 MP4.

	mask_frames: list of (H, W) bool arrays; True = person, False = background.
	Encodes via a temporary raw MP4 then re-encodes with ffmpeg libx264.
	"""
	if not mask_frames:
	return
	h, w = mask_frames[0].shape[:2]
	tmp = out_path + ".tmp.mp4"
	writer = cv2.VideoWriter(tmp, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h))
	for mask in mask_frames:
	gray = (mask * 255).astype(np.uint8)
	bgr = cv2.cvtColor(gray, cv2.COLOR_GRAY2BGR)
	writer.write(bgr)
	writer.release()
	ret = subprocess.run(
	["ffmpeg", "-y", "-i", tmp, "-c:v", "libx264", "-pix_fmt", "yuv420p", out_path],
	capture_output=True,
	)
	os.remove(tmp)
	if ret.returncode != 0:
	print(f" ffmpeg H.264 encode failed:\n{ret.stderr.decode()}")


	# ---------------------------------------------------------------------------
	# Core processing
	# ---------------------------------------------------------------------------

	def process_video(
	predictor,
	video_path: str,
	out_path: str,
	fps: float,
	) -> None:
	frames = load_video_frames(video_path)
	if not frames:
	print(f" Warning: no frames loaded from {video_path}")
	return

	# Start session
	response = predictor.handle_request(
	request=dict(type="start_session", resource_path=video_path)
	)
	session_id = response["session_id"]

	try:
	predictor.handle_request(
	request=dict(type="reset_session", session_id=session_id)
	)

	# Detect all person instances in frame 0
	response = predictor.handle_request(
	request=dict(
	type="add_prompt",
	session_id=session_id,
	frame_index=0,
	text="person",
	)
	)
	out0 = response["outputs"]
	obj_ids_0 = out0["out_obj_ids"].tolist()

	if not obj_ids_0:
	print(f" No persons detected — skipping")
	return

	print(f" Detected {len(obj_ids_0)} person(s): obj_ids={[int(x) for x in obj_ids_0]}")

	# Propagate through all frames — drain the stream before any numpy work
	# to avoid distributed-barrier deadlocks on multi-GPU setups.
	raw_outputs: dict[int, object] = {}
	for resp in predictor.handle_stream_request(
	request=dict(type="propagate_in_video", session_id=session_id)
	):
	frame_idx = resp["frame_index"]
	if frame_idx < len(frames):
	raw_outputs[frame_idx] = resp["outputs"]

	finally:
	predictor.handle_request(
	request=dict(type="close_session", session_id=session_id)
	)

	# Build one combined mask per frame (OR across all detected persons).
	# Frames with no tracking output get a black (all-zero) mask.
	n = len(frames)
	mask_frames: list[np.ndarray] = []
	for frame_idx in range(n):
	h, w = frames[frame_idx].shape[:2]
	combined = np.zeros((h, w), dtype=bool)
	if frame_idx in raw_outputs:
	out = raw_outputs[frame_idx]
	for i in range(len(out["out_obj_ids"].tolist())):
	combined \|= mask_to_numpy(out["out_binary_masks"][i])
	mask_frames.append(combined)

	save_mask_video(mask_frames, out_path, fps)
	print(f" Saved mask video ({n} frames) → {out_path}")


	# ---------------------------------------------------------------------------
	# Entry point
	# ---------------------------------------------------------------------------

	def main() -> None:
	parser = argparse.ArgumentParser(
	description="Extract person mask videos using SAM3. "
	"Single-file mode: --input / --output. "
	"Batch mode: --video_dir / --output_dir (and optionally --csv)."
	)
	# Single-file mode
	parser.add_argument("--input", default=None, help="Input .mp4 video file")
	parser.add_argument("--output", default=None, help="Output mask .mp4 file")
	# Batch mode
	parser.add_argument("--video_dir", default=None, help="Folder containing .mp4 files")
	parser.add_argument(
	"--csv", default=None,
	help="Optional CSV with a 'clip_id' column; if omitted, all .mp4s in "
	"video_dir are processed",
	)
	parser.add_argument(
	"--output_dir", default=None,
	help="Root directory for mask outputs (one masks.mp4 per clip subdirectory)",
	)
	parser.add_argument(
	"--checkpoint", default=None,
	help="Path to SAM 3 checkpoint .pt file (default: sam3/sam3/sam3.pt, "
	"or auto-download from HuggingFace)",
	)
	args = parser.parse_args()

	# Resolve mode and build (video_path, out_path) pairs
	entries: list[tuple[str, str]] = [] # (video_path, out_path)

	if args.input:
	if not args.output:
	parser.error("--input requires --output")
	entries = [(args.input, args.output)]
	elif args.video_dir:
	if not args.output_dir:
	parser.error("--video_dir requires --output_dir")
	if args.csv:
	with open(args.csv, newline="", encoding="utf-8") as f:
	reader = csv.DictReader(f)
	for row in reader:
	clip_id = row["clip_id"].strip()
	video_path = os.path.join(args.video_dir, clip_id + ".mp4")
	out_path = os.path.join(args.output_dir, clip_id, "masks.mp4")
	entries.append((video_path, out_path))
	print(f"Loaded {len(entries)} entries from {args.csv}")
	else:
	mp4s = sorted(glob.glob(os.path.join(args.video_dir, "*.mp4")))
	for p in mp4s:
	clip_id = Path(p).stem
	out_path = os.path.join(args.output_dir, clip_id, "masks.mp4")
	entries.append((p, out_path))
	print(f"Found {len(entries)} .mp4 file(s) in {args.video_dir}")
	else:
	parser.error("Provide either --input/--output or --video_dir/--output_dir")

	if not entries:
	print("No videos to process — exiting.")
	sys.exit(0)

	# Load model
	print("Loading SAM 3 model …")
	import torch
	from sam3.model_builder import build_sam3_video_predictor

	checkpoint = args.checkpoint
	if checkpoint is None:
	script_dir = os.path.dirname(os.path.abspath(__file__))
	local = os.path.join(script_dir, "sam3", "sam3", "sam3.pt")
	checkpoint = local if os.path.exists(local) else None

	gpus_to_use = list(range(torch.cuda.device_count())) or [torch.cuda.current_device()]
	predictor = build_sam3_video_predictor(checkpoint_path=checkpoint, gpus_to_use=gpus_to_use)
	print("Model ready.\n")

	# Process
	for video_path, out_path in entries:
	if not os.path.exists(video_path):
	print(f"[SKIP] Not found: {video_path}")
	continue
	os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
	fps = get_video_fps(video_path)
	print(f"[{Path(video_path).stem}] fps={fps:.1f} {video_path} → {out_path}")
	process_video(predictor=predictor, video_path=video_path, out_path=out_path, fps=fps)

	print("\nAll done.")


	if __name__ == "__main__":
	main()