Create app.py

app.py

@@ -0,0 +1,509 @@
+import json
+import logging
+import os
+import shutil
+import subprocess
+import sys
+import tempfile
+import traceback
+import uuid
+import zipfile
+from pathlib import Path
+
+import gradio as gr
+import spaces
+from huggingface_hub import hf_hub_download
+from PIL import Image
+
+
+logging.basicConfig(level=logging.INFO, format="[%(asctime)s] %(levelname)s: %(message)s")
+
+ROOT = Path(__file__).resolve().parent
+JOB_ROOT = Path(os.getenv("SCAIL_POSE_JOB_ROOT", str(Path(tempfile.gettempdir()) / "scail_pose_jobs")))
+JOB_ROOT.mkdir(parents=True, exist_ok=True)
+
+SAM3_REPO_ID = os.getenv("SCAIL_POSE_SAM3_REPO_ID", "facebook/sam3")
+SAM3_FILENAME = os.getenv("SCAIL_POSE_SAM3_FILENAME", "sam3.pt")
+WEIGHTS_DIR = Path(os.getenv("SCAIL_POSE_WEIGHTS_DIR", str(ROOT / "pretrained_weights")))
+SAM3_MODEL_PATH = Path(os.getenv("SCAIL_POSE_SAM3_MODEL", str(WEIGHTS_DIR / SAM3_FILENAME)))
+AUTO_DOWNLOAD_SAM3 = os.getenv("SCAIL_POSE_AUTO_DOWNLOAD_SAM3", "1") == "1"
+
+GPU_SIZE = os.getenv("SCAIL_POSE_ZEROGPU_SIZE", "xlarge")
+GPU_DURATION = int(os.getenv("SCAIL_POSE_GPU_DURATION", "600"))
+
+IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
+VIDEO_EXTS = {".mp4", ".mov", ".webm", ".mkv"}
+
+
+def _repo_status() -> str:
+    required = [
+        "NLFPoseExtract/process_animation_aio.py",
+        "NLFPoseExtract/process_replacement.py",
+        "NLFPoseExtract/v2_helper.py",
+        "TrackSam3/track.py",
+    ]
+    missing = [rel for rel in required if not (ROOT / rel).exists()]
+    if missing:
+        return (
+            "SCAIL-Pose repo layout was not found. Put this app.py at the root of "
+            "the SCAIL-Pose checkout.\n\nMissing:\n" + "\n".join(f"- {rel}" for rel in missing)
+        )
+
+    weight_status = (
+        f"SAM3 weights found: {SAM3_MODEL_PATH}"
+        if SAM3_MODEL_PATH.exists()
+        else f"SAM3 weights not found yet: {SAM3_MODEL_PATH}"
+    )
+    return (
+        "Ready. SCAIL-Pose repo layout detected.\n\n"
+        f"Job root: {JOB_ROOT}\n"
+        f"{weight_status}\n\n"
+        "This Space exports SCAIL-2-compatible input packs. Outputs are temporary."
+    )
+
+
+def _require_repo_layout():
+    missing = []
+    for rel in (
+        "NLFPoseExtract/process_animation_aio.py",
+        "NLFPoseExtract/process_replacement.py",
+        "NLFPoseExtract/v2_helper.py",
+        "TrackSam3/track.py",
+    ):
+        if not (ROOT / rel).exists():
+            missing.append(rel)
+    if missing:
+        raise RuntimeError(
+            "This app.py must live at the root of the SCAIL-Pose repository. "
+            f"Missing: {', '.join(missing)}"
+        )
+
+
+def _ensure_sam3_weights() -> Path:
+    if SAM3_MODEL_PATH.exists():
+        return SAM3_MODEL_PATH
+
+    if not AUTO_DOWNLOAD_SAM3:
+        raise RuntimeError(
+            f"SAM3 weights were not found at {SAM3_MODEL_PATH}. "
+            "Set SCAIL_POSE_SAM3_MODEL or enable SCAIL_POSE_AUTO_DOWNLOAD_SAM3=1."
+        )
+
+    WEIGHTS_DIR.mkdir(parents=True, exist_ok=True)
+    logging.info("Downloading SAM3 weights from %s/%s", SAM3_REPO_ID, SAM3_FILENAME)
+    downloaded = hf_hub_download(
+        repo_id=SAM3_REPO_ID,
+        filename=SAM3_FILENAME,
+        local_dir=str(WEIGHTS_DIR),
+        token=os.getenv("HF_TOKEN") or None,
+    )
+    return Path(downloaded)
+
+
+def _new_job_dir(mode: str) -> Path:
+    job_dir = JOB_ROOT / f"{mode}_{uuid.uuid4().hex}"
+    job_dir.mkdir(parents=True, exist_ok=False)
+    return job_dir
+
+
+def _as_path(upload, label: str) -> Path:
+    if upload is None:
+        raise RuntimeError(f"Missing {label}.")
+
+    if isinstance(upload, dict):
+        upload = upload.get("path") or upload.get("name")
+    elif hasattr(upload, "path"):
+        upload = upload.path
+    elif hasattr(upload, "name") and not isinstance(upload, (str, os.PathLike)):
+        upload = upload.name
+
+    path = Path(upload)
+    if not path.exists():
+        raise RuntimeError(f"{label} does not exist: {path}")
+    return path
+
+
+def _save_reference_image(upload, dest: Path) -> Path:
+    source = _as_path(upload, "reference image")
+    try:
+        image = Image.open(source).convert("RGB")
+        image.save(dest)
+    except Exception as exc:
+        raise RuntimeError(f"Could not read reference image: {source}") from exc
+    return dest
+
+
+def _copy_video(upload, dest: Path, label: str) -> Path:
+    source = _as_path(upload, label)
+    if source.suffix.lower() not in VIDEO_EXTS:
+        raise RuntimeError(f"{label} should be a video file. Got: {source.name}")
+    shutil.copy2(source, dest)
+    return dest
+
+
+def _text_args(text_prompt: str) -> list[str]:
+    words = [part.strip() for part in (text_prompt or "").split() if part.strip()]
+    return words or ["human", "character"]
+
+
+def _run_command(command: list[str], progress=None) -> str:
+    logging.info("Running command: %s", " ".join(command))
+    output_lines = []
+    proc = subprocess.Popen(
+        command,
+        cwd=str(ROOT),
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        bufsize=1,
+        env=os.environ.copy(),
+    )
+    assert proc.stdout is not None
+    for line in proc.stdout:
+        line = line.rstrip()
+        if not line:
+            continue
+        logging.info("[SCAIL-Pose] %s", line)
+        output_lines.append(line)
+        if progress is not None:
+            progress(None, desc=line[:120])
+
+    ret = proc.wait()
+    output = "\n".join(output_lines)
+    if ret != 0:
+        raise RuntimeError(f"SCAIL-Pose command failed with exit code {ret}.\n\n{output}")
+    return output
+
+
+def _write_metadata(job_dir: Path, mode: str, prompt: str, extra: dict | None = None) -> None:
+    metadata = {
+        "mode": mode,
+        "source": "scail-pose-gradio-pack-builder",
+    }
+    if extra:
+        metadata.update(extra)
+    (job_dir / "metadata.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8")
+    (job_dir / "prompt.txt").write_text(prompt or "", encoding="utf-8")
+
+
+def _zip_pack(job_dir: Path, mode: str) -> Path:
+    zip_path = job_dir / f"scail2_{mode}_pack.zip"
+    include = [
+        "ref.png",
+        "ref_mask.png",
+        "ref_mask.jpg",
+        "rendered_v2.mp4",
+        "rendered_mask_v2.mp4",
+        "replace_mask.mp4",
+        "prompt.txt",
+        "metadata.json",
+    ]
+    with zipfile.ZipFile(zip_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+        for name in include:
+            path = job_dir / name
+            if path.exists():
+                zf.write(path, arcname=name)
+    return zip_path
+
+
+def _require_outputs(job_dir: Path, names: list[str]) -> None:
+    missing = [name for name in names if not (job_dir / name).exists()]
+    if missing:
+        raise RuntimeError("SCAIL-Pose did not produce expected output(s): " + ", ".join(missing))
+
+
+def _animation_command(
+    job_dir: Path,
+    sam3_model: Path,
+    max_persons: int,
+    text_prompt: str,
+    crop_mode: str,
+) -> list[str]:
+    command = [
+        sys.executable,
+        str(ROOT / "NLFPoseExtract" / "process_animation_aio.py"),
+        "--subdir",
+        str(job_dir),
+        "--video_name",
+        "driving.mp4",
+        "--e2e_mode",
+        "--max_persons",
+        str(int(max_persons)),
+        "--sam3_model",
+        str(sam3_model),
+        "--text",
+        *_text_args(text_prompt),
+    ]
+    if crop_mode == "mask silhouette":
+        command.append("--crop_e2e_mask")
+    elif crop_mode == "moving bbox":
+        command.append("--crop_e2e_bbox")
+    elif crop_mode == "steady bbox":
+        command.append("--crop_e2e_steady_bbox")
+    return command
+
+
+def _replacement_command(
+    job_dir: Path,
+    sam3_model: Path,
+    text_prompt: str,
+    matchnearest: bool,
+    egocentric: bool,
+) -> list[str]:
+    command = [
+        sys.executable,
+        str(ROOT / "NLFPoseExtract" / "process_replacement.py"),
+        "--subdir",
+        str(job_dir),
+        "--video_name",
+        "driving.mp4",
+        "--sam3_model",
+        str(sam3_model),
+        "--text",
+        *_text_args(text_prompt),
+    ]
+    if matchnearest:
+        command.append("--matchnearest")
+    if egocentric:
+        command.append("--egocentric")
+    return command
+
+
+@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
+def build_animation_pack(
+    ref_image,
+    driving_video,
+    prompt,
+    sam3_text,
+    max_persons,
+    crop_mode,
+    progress=gr.Progress(track_tqdm=True),
+):
+    try:
+        progress(0.0, desc="Checking SCAIL-Pose repo")
+        _require_repo_layout()
+        progress(0.04, desc="Preparing SAM3 weights")
+        sam3_model = _ensure_sam3_weights()
+
+        job_dir = _new_job_dir("animation")
+        progress(0.08, desc="Preparing inputs")
+        _save_reference_image(ref_image, job_dir / "ref.png")
+        _copy_video(driving_video, job_dir / "driving.mp4", "driving video")
+        _write_metadata(
+            job_dir,
+            "animation",
+            prompt,
+            {
+                "driving": {
+                    "video": "rendered_v2.mp4",
+                    "mask_video": "rendered_mask_v2.mp4",
+                },
+                "primary": {
+                    "image": "ref.png",
+                    "mask": "ref_mask.jpg",
+                },
+                "sam3_text": _text_args(sam3_text),
+                "max_persons": int(max_persons),
+                "crop_mode": crop_mode,
+            },
+        )
+
+        progress(0.12, desc="Generating animation masks")
+        logs = _run_command(
+            _animation_command(job_dir, sam3_model, int(max_persons), sam3_text, crop_mode),
+            progress=progress,
+        )
+        _require_outputs(job_dir, ["ref_mask.jpg", "rendered_v2.mp4", "rendered_mask_v2.mp4"])
+
+        progress(0.92, desc="Packaging SCAIL-2 input pack")
+        zip_path = _zip_pack(job_dir, "animation")
+        progress(1.0, desc="Done")
+        status = f"Done. Pack created at {zip_path}\n\n{logs}"
+        return (
+            str(job_dir / "ref_mask.jpg"),
+            str(job_dir / "rendered_v2.mp4"),
+            str(job_dir / "rendered_mask_v2.mp4"),
+            str(zip_path),
+            status,
+        )
+    except Exception:
+        logging.exception("Animation pack generation failed")
+        return None, None, None, None, traceback.format_exc()
+
+
+@spaces.GPU(duration=GPU_DURATION, size=GPU_SIZE)
+def build_replacement_pack(
+    ref_image,
+    driving_video,
+    prompt,
+    sam3_text,
+    matchnearest,
+    egocentric,
+    progress=gr.Progress(track_tqdm=True),
+):
+    try:
+        if matchnearest and egocentric:
+            raise RuntimeError("matchnearest and egocentric are mutually exclusive.")
+
+        progress(0.0, desc="Checking SCAIL-Pose repo")
+        _require_repo_layout()
+        progress(0.04, desc="Preparing SAM3 weights")
+        sam3_model = _ensure_sam3_weights()
+
+        job_dir = _new_job_dir("replacement")
+        progress(0.08, desc="Preparing inputs")
+        _save_reference_image(ref_image, job_dir / "ref.png")
+        _copy_video(driving_video, job_dir / "driving.mp4", "source video")
+        _write_metadata(
+            job_dir,
+            "replacement",
+            prompt,
+            {
+                "driving": {
+                    "video": "rendered_v2.mp4",
+                    "mask_video": "replace_mask.mp4",
+                },
+                "primary": {
+                    "image": "ref.png",
+                    "mask": "ref_mask.png",
+                },
+                "sam3_text": _text_args(sam3_text),
+                "matchnearest": bool(matchnearest),
+                "egocentric": bool(egocentric),
+            },
+        )
+
+        progress(0.12, desc="Generating replacement masks")
+        logs = _run_command(
+            _replacement_command(job_dir, sam3_model, sam3_text, bool(matchnearest), bool(egocentric)),
+            progress=progress,
+        )
+        _require_outputs(job_dir, ["ref_mask.png", "rendered_v2.mp4", "replace_mask.mp4"])
+
+        progress(0.92, desc="Packaging SCAIL-2 input pack")
+        zip_path = _zip_pack(job_dir, "replacement")
+        progress(1.0, desc="Done")
+        status = f"Done. Pack created at {zip_path}\n\n{logs}"
+        return (
+            str(job_dir / "ref_mask.png"),
+            str(job_dir / "rendered_v2.mp4"),
+            str(job_dir / "replace_mask.mp4"),
+            str(zip_path),
+            status,
+        )
+    except Exception:
+        logging.exception("Replacement pack generation failed")
+        return None, None, None, None, traceback.format_exc()
+
+
+def build_ui():
+    with gr.Blocks(title="SCAIL-Pose Pack Builder") as demo:
+        gr.Markdown(
+            "# SCAIL-Pose Pack Builder\n"
+            "Generate SCAIL-2-ready masks and export them as an input pack. "
+            "Use the downloaded `.zip` in the SCAIL-2 demo Advanced Pack tab."
+        )
+        gr.Textbox(value=_repo_status(), label="Startup status", interactive=False, lines=7)
+
+        with gr.Tab("Animation Pack"):
+            gr.Markdown(
+                "Create an animation pack from one reference image and one driving video. "
+                "This uses SCAIL-Pose end-to-end mode: the driving video becomes `rendered_v2.mp4`, "
+                "and SAM3 produces the colored driving mask video."
+            )
+            with gr.Row():
+                anim_ref = gr.Image(type="filepath", label="Reference image")
+                anim_driving = gr.Video(label="Driving video")
+            anim_prompt = gr.Textbox(label="Prompt for SCAIL-2", lines=3)
+            with gr.Row():
+                anim_sam3_text = gr.Textbox(value="human character", label="SAM3 text prompt")
+                anim_max_persons = gr.Number(value=2, precision=0, label="Max tracked subjects")
+                anim_crop = gr.Dropdown(
+                    ["none", "mask silhouette", "moving bbox", "steady bbox"],
+                    value="none",
+                    label="Driving crop mode",
+                )
+            anim_run = gr.Button("Generate animation pack", variant="primary")
+            with gr.Row():
+                anim_ref_mask = gr.Image(label="Reference mask", interactive=False)
+                anim_mask_video = gr.Video(label="Driving mask video")
+            anim_rendered = gr.Video(label="Rendered / driving video")
+            anim_zip = gr.File(label="Download SCAIL-2 animation pack")
+            anim_status = gr.Textbox(label="Run logs", lines=14)
+            anim_run.click(
+                build_animation_pack,
+                inputs=[
+                    anim_ref,
+                    anim_driving,
+                    anim_prompt,
+                    anim_sam3_text,
+                    anim_max_persons,
+                    anim_crop,
+                ],
+                outputs=[anim_ref_mask, anim_rendered, anim_mask_video, anim_zip, anim_status],
+            )
+
+        with gr.Tab("Replacement Pack"):
+            gr.Markdown(
+                "Create a replacement pack from a replacement reference image and a source video. "
+                "The source video becomes `rendered_v2.mp4`, and SAM3 produces `replace_mask.mp4`."
+            )
+            with gr.Row():
+                repl_ref = gr.Image(type="filepath", label="Replacement reference image")
+                repl_driving = gr.Video(label="Source / driving video")
+            repl_prompt = gr.Textbox(label="Prompt for SCAIL-2", lines=3)
+            with gr.Row():
+                repl_sam3_text = gr.Textbox(value="human character", label="SAM3 text prompt")
+                repl_matchnearest = gr.Checkbox(value=False, label="Match nearest target when two people are detected")
+                repl_egocentric = gr.Checkbox(value=False, label="Egocentric: union two detected parts into one actor")
+            repl_run = gr.Button("Generate replacement pack", variant="primary")
+            with gr.Row():
+                repl_ref_mask = gr.Image(label="Reference mask", interactive=False)
+                repl_mask_video = gr.Video(label="Replacement mask video")
+            repl_rendered = gr.Video(label="Rendered / source video")
+            repl_zip = gr.File(label="Download SCAIL-2 replacement pack")
+            repl_status = gr.Textbox(label="Run logs", lines=14)
+            repl_run.click(
+                build_replacement_pack,
+                inputs=[
+                    repl_ref,
+                    repl_driving,
+                    repl_prompt,
+                    repl_sam3_text,
+                    repl_matchnearest,
+                    repl_egocentric,
+                ],
+                outputs=[repl_ref_mask, repl_rendered, repl_mask_video, repl_zip, repl_status],
+            )
+
+        with gr.Tab("Pack Format"):
+            gr.Markdown(
+                "The exported zip is intentionally flat for V1, so it is easy to inspect and "
+                "compatible with the SCAIL-2 Advanced Pack parser.\n\n"
+                "Animation pack:\n"
+                "```text\n"
+                "ref.png\n"
+                "ref_mask.jpg\n"
+                "rendered_v2.mp4\n"
+                "rendered_mask_v2.mp4\n"
+                "prompt.txt\n"
+                "metadata.json\n"
+                "```\n\n"
+                "Replacement pack:\n"
+                "```text\n"
+                "ref.png\n"
+                "ref_mask.png\n"
+                "rendered_v2.mp4\n"
+                "replace_mask.mp4\n"
+                "prompt.txt\n"
+                "metadata.json\n"
+                "```\n\n"
+                "For multi-character curation, run several focused jobs first, inspect the masks, "
+                "then combine the selected references into a canonical Advanced Pack."
+            )
+
+    return demo
+
+
+if __name__ == "__main__":
+    build_ui().queue(max_size=4).launch(show_error=True)