scripts/build_hyperview_demo_assets.py

#!/usr/bin/env python3
"""Build precomputed HyperView embedding assets for the jaguar Space."""

from __future__ import annotations

import argparse
import json
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
import sys
from typing import Any
from urllib.parse import urlparse

import numpy as np
import pandas as pd
import torch
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

PROJECT_ROOT = Path(__file__).resolve().parents[2]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.append(str(PROJECT_ROOT))

from experiment_scripts.evaluate_inpainted_bgfg import (  # noqa: E402
    _load_arcface_benchmark,
    _load_lorentz,
    _load_triplet_benchmark,
)
from experiment_scripts.train_lorentz_reid import build_transforms  # noqa: E402

DEFAULT_MANIFEST_PATH = PROJECT_ROOT / "HyperViewDemoHuggingFaceSpace/config/model_manifest.json"
DEFAULT_DATASET_ROOT = PROJECT_ROOT / "kaggle_jaguar_dataset_v2"
DEFAULT_CORESET_CSV = PROJECT_ROOT / "data/validation_coreset.csv"
DEFAULT_OUTPUT_DIR = PROJECT_ROOT / "HyperViewDemoHuggingFaceSpace/assets"


@dataclass
class LoadedModel:
    model: Any
    val_transform: Any
    image_size: int


class JaguarEmbeddingDataset(Dataset):
    def __init__(
        self,
        rows: list[dict[str, str]],
        images_dir: Path,
        transform: Any,
        image_variant: str,
    ):
        self.rows = rows
        self.images_dir = images_dir
        self.transform = transform
        self.image_variant = image_variant

    def __len__(self) -> int:
        return len(self.rows)

    @staticmethod
    def _is_albumentations_transform(transform: Any) -> bool:
        return transform.__class__.__module__.startswith("albumentations")

    def _load_image(self, filename: str) -> Image.Image:
        image_path = self.images_dir / filename
        if self.image_variant == "foreground_only":
            rgba = Image.open(image_path).convert("RGBA")
            rgba_np = np.array(rgba, dtype=np.uint8)
            rgb = rgba_np[:, :, :3]
            alpha = rgba_np[:, :, 3]
            mask = (alpha > 0).astype(np.uint8)
            cutout_rgb = (rgb * mask[:, :, np.newaxis]).astype(np.uint8)
            return Image.fromarray(cutout_rgb, mode="RGB")
        return Image.open(image_path).convert("RGB")

    def __getitem__(self, idx: int):
        row = self.rows[idx]
        image = self._load_image(row["filename"])

        if self.transform is None:
            raise ValueError("Validation transform is required for embedding extraction.")

        if self._is_albumentations_transform(self.transform):
            image_tensor = self.transform(image=np.array(image, dtype=np.uint8))["image"]
        else:
            image_tensor = self.transform(image)

        return (
            image_tensor,
            row["sample_id"],
            row["label"],
            row["filename"],
            row["split_tag"],
        )


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Build precomputed embedding artifacts for HyperView Space runtime."
    )
    parser.add_argument(
        "--model_manifest",
        type=Path,
        default=DEFAULT_MANIFEST_PATH,
        help="Model manifest JSON defining the three demo models.",
    )
    parser.add_argument(
        "--dataset_root",
        type=Path,
        default=DEFAULT_DATASET_ROOT,
        help="Dataset root containing train.csv and train/ images.",
    )
    parser.add_argument(
        "--coreset_csv",
        type=Path,
        default=DEFAULT_CORESET_CSV,
        help="Validation coreset CSV used to tag split_tag=train/validation.",
    )
    parser.add_argument(
        "--output_dir",
        type=Path,
        default=DEFAULT_OUTPUT_DIR,
        help="Output directory for per-model embeddings and manifest JSON.",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        choices=["cuda"],
        help="Runtime device. CUDA-only by contract.",
    )
    parser.add_argument("--batch_size", type=int, default=64)
    parser.add_argument("--num_workers", type=int, default=4)
    parser.add_argument(
        "--image_variant",
        type=str,
        default="foreground_only",
        choices=["foreground_only", "full_rgb"],
    )
    parser.add_argument(
        "--max_samples",
        type=int,
        default=None,
        help="Optional smoke-mode sample cap for quick checks.",
    )
    return parser.parse_args()


def utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def resolve_device(device_name: str) -> torch.device:
    if device_name != "cuda":
        raise SystemExit("GPU unavailable: CUDA requested but not available.")
    if not torch.cuda.is_available():
        raise SystemExit("GPU unavailable: CUDA requested but not available.")
    return torch.device("cuda")


def load_model_manifest(manifest_path: Path) -> dict[str, Any]:
    payload = json.loads(manifest_path.read_text(encoding="utf-8"))
    if "models" not in payload or not isinstance(payload["models"], list):
        raise ValueError(f"Invalid model manifest: {manifest_path}")
    return payload


def parse_run_url(run_url: str) -> tuple[str, str, str]:
    parsed = urlparse(run_url)
    parts = [p for p in parsed.path.split("/") if p]
    if len(parts) >= 4 and parts[2] == "runs":
        return parts[0], parts[1], parts[3]
    raise ValueError(f"Unsupported W&B run URL format: {run_url}")


def pick_checkpoint_file(root: Path, checkpoint_name: str | None) -> Path:
    if checkpoint_name:
        exact = sorted(root.rglob(checkpoint_name))
        if exact:
            return exact[0]

    candidates = sorted(root.rglob("*.pth"))
    if not candidates:
        raise FileNotFoundError(f"No .pth checkpoints found under downloaded artifact: {root}")
    return candidates[0]


def download_checkpoint_from_wandb(
    run_url: str,
    model_key: str,
    checkpoint_name: str | None,
    output_dir: Path,
) -> tuple[Path, str]:
    try:
        import wandb
    except ImportError as exc:
        raise ImportError(
            "wandb is required to download missing checkpoints. Install with `uv pip install wandb`."
        ) from exc

    entity, project, run_id = parse_run_url(run_url)
    api = wandb.Api()
    run = api.run(f"{entity}/{project}/{run_id}")

    artifacts = [artifact for artifact in run.logged_artifacts() if artifact.type == "model"]
    if not artifacts:
        raise FileNotFoundError(
            f"No model artifacts found for run {entity}/{project}/{run_id}."
        )

    artifact = artifacts[-1]
    safe_name = artifact.name.replace("/", "_").replace(":", "_")
    download_root = output_dir / "downloaded_checkpoints" / model_key / safe_name
    download_root.mkdir(parents=True, exist_ok=True)
    downloaded_dir = Path(artifact.download(root=str(download_root)))

    checkpoint_path = pick_checkpoint_file(downloaded_dir, checkpoint_name)
    return checkpoint_path, f"wandb_artifact:{artifact.name}"


def resolve_checkpoint_path(model_cfg: dict[str, Any], output_dir: Path) -> tuple[Path, str]:
    checkpoint_path = Path(model_cfg.get("checkpoint_path", ""))
    if not checkpoint_path.is_absolute():
        checkpoint_path = (PROJECT_ROOT / checkpoint_path).resolve()

    if checkpoint_path.exists():
        return checkpoint_path, "local_path"

    run_url = model_cfg.get("run_url")
    if not run_url:
        raise FileNotFoundError(
            f"Checkpoint not found at {checkpoint_path} and no run_url provided for fallback download."
        )

    return download_checkpoint_from_wandb(
        run_url=run_url,
        model_key=str(model_cfg["model_key"]),
        checkpoint_name=model_cfg.get("checkpoint_name"),
        output_dir=output_dir,
    )


def read_augmentation_profile(checkpoint_path: Path) -> str:
    checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
    return str(checkpoint.get("augmentation_profile", "lorentz_default"))


def load_model(model_cfg: dict[str, Any], checkpoint_path: Path, device: str) -> LoadedModel:
    loader = str(model_cfg["loader"])
    if loader == "arcface_benchmark":
        model, image_size, _metric = _load_arcface_benchmark(str(checkpoint_path), device)
        augmentation_profile = read_augmentation_profile(checkpoint_path)
        _train_tf, val_tf, _resolved = build_transforms(image_size, augmentation_profile=augmentation_profile)
        return LoadedModel(model=model, val_transform=val_tf, image_size=int(image_size))

    if loader == "triplet_benchmark":
        model, image_size, _metric = _load_triplet_benchmark(str(checkpoint_path), device)
        augmentation_profile = read_augmentation_profile(checkpoint_path)
        _train_tf, val_tf, _resolved = build_transforms(image_size, augmentation_profile=augmentation_profile)
        return LoadedModel(model=model, val_transform=val_tf, image_size=int(image_size))

    if loader == "lorentz":
        model, image_size, _metric, val_tf = _load_lorentz(str(checkpoint_path), device)
        return LoadedModel(model=model, val_transform=val_tf, image_size=int(image_size))

    raise ValueError(f"Unsupported loader='{loader}' in model manifest.")


def build_sample_rows(
    dataset_root: Path,
    coreset_csv: Path,
    max_samples: int | None,
) -> list[dict[str, str]]:
    train_csv = dataset_root / "train.csv"
    images_dir = dataset_root / "train"
    if not train_csv.exists():
        raise FileNotFoundError(f"Missing train.csv at {train_csv}")
    if not images_dir.exists():
        raise FileNotFoundError(f"Missing train images directory at {images_dir}")

    train_df = pd.read_csv(train_csv)
    coreset_df = pd.read_csv(coreset_csv)
    coreset_filenames = set(coreset_df["filename"].astype(str).tolist())

    train_df = train_df.copy()
    train_df["filename"] = train_df["filename"].astype(str)
    train_df["ground_truth"] = train_df["ground_truth"].astype(str)
    train_df["sample_id"] = train_df["filename"]
    train_df["split_tag"] = np.where(train_df["filename"].isin(coreset_filenames), "validation", "train")

    if max_samples is not None:
        train_df = train_df.iloc[: int(max_samples)].copy()

    rows: list[dict[str, str]] = []
    for _, row in train_df.iterrows():
        rows.append(
            {
                "sample_id": str(row["sample_id"]),
                "filename": str(row["filename"]),
                "label": str(row["ground_truth"]),
                "split_tag": str(row["split_tag"]),
            }
        )

    return rows


def extract_embeddings(
    loaded_model: LoadedModel,
    rows: list[dict[str, str]],
    images_dir: Path,
    image_variant: str,
    device: torch.device,
    batch_size: int,
    num_workers: int,
    progress_label: str,
) -> tuple[list[str], np.ndarray, list[str], list[str], list[str]]:
    dataset = JaguarEmbeddingDataset(
        rows=rows,
        images_dir=images_dir,
        transform=loaded_model.val_transform,
        image_variant=image_variant,
    )
    loader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=False,
        num_workers=num_workers,
        pin_memory=True,
    )

    all_vectors: list[np.ndarray] = []
    all_ids: list[str] = []
    all_labels: list[str] = []
    all_filenames: list[str] = []
    all_split_tags: list[str] = []

    loaded_model.model.eval()
    with torch.no_grad():
        for images, sample_ids, labels, filenames, split_tags in tqdm(loader, desc=progress_label):
            images = images.to(device, non_blocking=True)
            vectors = loaded_model.model(images)
            if isinstance(vectors, (tuple, list)):
                vectors = vectors[0]
            vectors_np = vectors.detach().cpu().numpy().astype(np.float32)

            all_vectors.append(vectors_np)
            all_ids.extend([str(x) for x in sample_ids])
            all_labels.extend([str(x) for x in labels])
            all_filenames.extend([str(x) for x in filenames])
            all_split_tags.extend([str(x) for x in split_tags])

    if not all_vectors:
        raise RuntimeError("No embeddings were generated.")

    stacked = np.vstack(all_vectors).astype(np.float32)
    return all_ids, stacked, all_labels, all_filenames, all_split_tags


def save_model_artifacts(
    output_dir: Path,
    model_cfg: dict[str, Any],
    checkpoint_path: Path,
    checkpoint_source: str,
    sample_ids: list[str],
    vectors: np.ndarray,
    labels: list[str],
    filenames: list[str],
    split_tags: list[str],
    image_variant: str,
    image_size: int,
    batch_size: int,
    num_workers: int,
) -> dict[str, Any]:
    model_key = str(model_cfg["model_key"])
    model_dir = output_dir / "models" / model_key
    model_dir.mkdir(parents=True, exist_ok=True)

    embeddings_path = model_dir / "embeddings.npz"
    metadata_path = model_dir / "metadata.json"

    np.savez_compressed(
        embeddings_path,
        ids=np.asarray(sample_ids),
        vectors=vectors,
        labels=np.asarray(labels),
        filenames=np.asarray(filenames),
        split_tags=np.asarray(split_tags),
    )

    metadata = {
        "generated_at_utc": utc_now(),
        "model_key": model_key,
        "comparison_key": model_cfg.get("comparison_key"),
        "family": model_cfg.get("family"),
        "loader": model_cfg.get("loader"),
        "space_key": model_cfg.get("space_key"),
        "geometry": model_cfg.get("geometry"),
        "layout": model_cfg.get("layout"),
        "num_samples": int(vectors.shape[0]),
        "embedding_dim": int(vectors.shape[1]),
        "checkpoint_path": str(checkpoint_path),
        "checkpoint_source": checkpoint_source,
        "run_url": model_cfg.get("run_url"),
        "image_variant": image_variant,
        "image_size": int(image_size),
        "batch_size": int(batch_size),
        "num_workers": int(num_workers),
    }
    metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8")

    return {
        "model_key": model_key,
        "comparison_key": model_cfg.get("comparison_key"),
        "family": model_cfg.get("family"),
        "loader": model_cfg.get("loader"),
        "space_key": model_cfg.get("space_key"),
        "geometry": model_cfg.get("geometry"),
        "layout": model_cfg.get("layout"),
        "checkpoint_path": str(checkpoint_path),
        "checkpoint_source": checkpoint_source,
        "run_url": model_cfg.get("run_url"),
        "embeddings_path": str(embeddings_path.relative_to(output_dir)),
        "metadata_path": str(metadata_path.relative_to(output_dir)),
        "num_samples": int(vectors.shape[0]),
        "embedding_dim": int(vectors.shape[1]),
    }


def write_sample_index(output_dir: Path, rows: list[dict[str, str]]) -> Path:
    sample_index_path = output_dir / "sample_index.csv"
    sample_df = pd.DataFrame(rows)
    sample_df.to_csv(sample_index_path, index=False)
    return sample_index_path


def main() -> int:
    args = parse_args()
    device = resolve_device(args.device)

    model_manifest = load_model_manifest(args.model_manifest)
    output_dir = args.output_dir.resolve()
    output_dir.mkdir(parents=True, exist_ok=True)

    dataset_root = args.dataset_root.resolve()
    images_dir = dataset_root / "train"
    rows = build_sample_rows(
        dataset_root=dataset_root,
        coreset_csv=args.coreset_csv,
        max_samples=args.max_samples,
    )
    if not rows:
        raise RuntimeError("No rows found in train.csv after applying filters.")

    expected_ids = [row["sample_id"] for row in rows]
    sample_index_path = write_sample_index(output_dir, rows)

    emitted_models: list[dict[str, Any]] = []
    for model_cfg in model_manifest["models"]:
        model_key = str(model_cfg["model_key"])
        print(f"\n=== Building embeddings for {model_key} ===")

        checkpoint_path, checkpoint_source = resolve_checkpoint_path(model_cfg=model_cfg, output_dir=output_dir)
        print(f"Checkpoint: {checkpoint_path} ({checkpoint_source})")

        loaded_model = load_model(model_cfg=model_cfg, checkpoint_path=checkpoint_path, device=args.device)
        ids, vectors, labels, filenames, split_tags = extract_embeddings(
            loaded_model=loaded_model,
            rows=rows,
            images_dir=images_dir,
            image_variant=args.image_variant,
            device=device,
            batch_size=int(args.batch_size),
            num_workers=int(args.num_workers),
            progress_label=f"extract:{model_key}",
        )

        if ids != expected_ids:
            raise RuntimeError(
                f"Sample ID alignment failed for {model_key}: extracted order does not match expected sample index."
            )

        emitted = save_model_artifacts(
            output_dir=output_dir,
            model_cfg=model_cfg,
            checkpoint_path=checkpoint_path,
            checkpoint_source=checkpoint_source,
            sample_ids=ids,
            vectors=vectors,
            labels=labels,
            filenames=filenames,
            split_tags=split_tags,
            image_variant=args.image_variant,
            image_size=loaded_model.image_size,
            batch_size=int(args.batch_size),
            num_workers=int(args.num_workers),
        )
        emitted_models.append(emitted)

    manifest_out = {
        "generated_at_utc": utc_now(),
        "source_model_manifest": str(args.model_manifest.resolve()),
        "dataset": {
            "dataset_root": str(dataset_root),
            "images_dir": str(images_dir),
            "coreset_csv": str(args.coreset_csv.resolve()),
            "num_samples": len(rows),
            "image_variant": args.image_variant,
            "sample_index_csv": str(sample_index_path.relative_to(output_dir)),
        },
        "models": emitted_models,
    }

    manifest_path = output_dir / "manifest.json"
    manifest_path.write_text(json.dumps(manifest_out, indent=2), encoding="utf-8")

    print("\n=== HyperView asset build complete ===")
    print(f"Sample count: {len(rows)}")
    print(f"Manifest: {manifest_path}")
    for emitted in emitted_models:
        print(
            f"- {emitted['model_key']}: {emitted['num_samples']} x {emitted['embedding_dim']} "
            f"({emitted['embeddings_path']})"
        )

    return 0


if __name__ == "__main__":
    raise SystemExit(main())