code/VESSA/data_preprocessing/build_template_bank.py

#!/usr/bin/env python3
"""
Build a template bank (template/*.npz + template_summary.json) for a raw dataset,
matching the schema used by existing <name>_10percent_processed directories.

Inputs:
  A raw dataset dir of the form:
    <source_dir>/
      dataset_info.json         (must contain "dataset_name" and ideally "modality" and "classes"/"class_mapping")
      train/images/             (*.jpg | *.png | *.jpeg)
      train/masks/              (matching *.png masks)

Outputs:
  <output_dir>/
    template/
      {stem}.npz                one per chosen sample, with keys:
        - template_image  uint8 [H, W, 3]   (NATIVE image size, no resize)
        - template_mask   uint8 [H, W]      (multi-class, 0 = background)
        - cls_feature     float32 [384]     (DINOv2-vits14 CLS, cosine-normalized)
        - classes         int64  [K]        (non-bg class ids present in this mask)
        - original_size   int64  [W, H]
        - modality        unicode str       (from dataset_info.json, else "Unknown")
      template_summary.json     index of all templates with metadata

Selection:
  ratio × N_train samples are drawn uniformly at random from train/images, seeded by --seed.
  Each chosen sample is ALSO a valid SFT target sample (bank ⊂ train). No disjointness.
"""
import argparse
import json
import os
import random
import sys
from pathlib import Path

import cv2
import numpy as np
from PIL import Image


THIS_DIR = Path(__file__).resolve().parent
VESSA_ROOT = THIS_DIR.parent
TRAIN_VESSA = VESSA_ROOT / "train_vessa"
if str(TRAIN_VESSA) not in sys.path:
    sys.path.insert(0, str(TRAIN_VESSA))


IMAGE_EXTS = (".jpg", ".jpeg", ".png")
MASK_EXTS = (".png", ".jpg", ".jpeg")


def _list_stems(root: str, exts):
    if not os.path.isdir(root):
        return []
    stems = []
    for name in sorted(os.listdir(root)):
        low = name.lower()
        if any(low.endswith(e) for e in exts):
            stems.append(os.path.splitext(name)[0])
    return stems


def _path_for_stem(root: str, stem: str, exts):
    for name in os.listdir(root):
        base, ext = os.path.splitext(name)
        if base == stem and any(name.lower().endswith(e) for e in exts):
            return os.path.join(root, name)
    return None


def _load_dataset_info(source_dir: str):
    info_path = os.path.join(source_dir, "dataset_info.json")
    if not os.path.isfile(info_path):
        raise FileNotFoundError(f"dataset_info.json not found in {source_dir}")
    with open(info_path, "r", encoding="utf-8") as f:
        return json.load(f)


def _unique_fg_classes(mask: np.ndarray):
    return sorted(int(v) for v in np.unique(mask).tolist() if int(v) != 0)


def _choose_bank_stems(train_img_dir: str, train_mask_dir: str, ratio: float, seed: int):
    img_stems = set(_list_stems(train_img_dir, IMAGE_EXTS))
    msk_stems = set(_list_stems(train_mask_dir, MASK_EXTS))
    common = sorted(img_stems & msk_stems)
    n_total = len(common)
    if n_total == 0:
        raise RuntimeError(f"No (image, mask) pairs under {train_img_dir} / {train_mask_dir}")
    n_bank = max(1, int(round(n_total * ratio)))
    rng = random.Random(seed)
    chosen = rng.sample(common, n_bank)
    return common, chosen


def build_bank(source_dir: str, output_dir: str, ratio: float, seed: int,
               device: str = "cuda:0", model_name: str = "dinov2_vits14"):
    # Lazy-import so the module can load without GPU libs when just reading
    from template_manager import TemplateManager

    info = _load_dataset_info(source_dir)
    dataset_name = info.get("dataset_name", os.path.basename(source_dir.rstrip("/")))
    modality = info.get("modality", "Unknown")

    train_img_dir = os.path.join(source_dir, "train", "images")
    train_mask_dir = os.path.join(source_dir, "train", "masks")

    all_stems, chosen_stems = _choose_bank_stems(train_img_dir, train_mask_dir, ratio, seed)
    print(f"[build_bank] {dataset_name}: n_train={len(all_stems)} ratio={ratio:.0%} bank_size={len(chosen_stems)}")

    template_dir = os.path.join(output_dir, "template")
    os.makedirs(template_dir, exist_ok=True)

    # DINOv2 feature extractor (cosine-normalized, matches TemplateManager default)
    tm = TemplateManager(device=device, model_name=model_name, distance_metric="cosine")

    summary = {}
    built = 0
    skipped = 0
    for stem in chosen_stems:
        img_path = _path_for_stem(train_img_dir, stem, IMAGE_EXTS)
        msk_path = _path_for_stem(train_mask_dir, stem, MASK_EXTS)
        if img_path is None or msk_path is None:
            skipped += 1
            continue

        img_bgr = cv2.imread(img_path, cv2.IMREAD_COLOR)
        mask = cv2.imread(msk_path, cv2.IMREAD_GRAYSCALE)
        if img_bgr is None or mask is None:
            skipped += 1
            continue
        img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)

        classes = _unique_fg_classes(mask)
        if not classes:
            # A mask with no foreground is useless as a template.
            skipped += 1
            continue

        # DINOv2 cls_feature at native resolution (TemplateManager resizes to 224 internally)
        cls_feat = tm._extract_feature(Image.fromarray(img_rgb))

        H, W = img_rgb.shape[:2]
        out_path = os.path.join(template_dir, stem + ".npz")
        np.savez_compressed(
            out_path,
            template_image=img_rgb.astype(np.uint8),
            template_mask=mask.astype(np.uint8),
            cls_feature=cls_feat.astype(np.float32),
            classes=np.array(classes, dtype=np.int64),
            original_size=np.array([W, H], dtype=np.int64),
            modality=np.array(modality, dtype="U10"),
        )
        summary[stem] = {
            "path": out_path,
            "original_size": [W, H],
            "cls_feature_shape": [int(cls_feat.shape[0])],
            "classes": classes,
            "modality": modality,
        }
        built += 1

    summary_path = os.path.join(template_dir, "template_summary.json")
    with open(summary_path, "w", encoding="utf-8") as f:
        json.dump(summary, f, indent=2)

    print(f"[build_bank] {dataset_name}: built={built} skipped={skipped} -> {template_dir}")
    return built, skipped


def parse_args():
    p = argparse.ArgumentParser(description="Build a DINOv2-feature template bank at a given ratio.")
    p.add_argument("--source_dir", required=True, type=str,
                   help="Raw dataset root (contains dataset_info.json, train/, val/)")
    p.add_argument("--output_dir", required=True, type=str,
                   help="Output dataset root (e.g. <name>_5percent_processed). template/ will be created under this.")
    p.add_argument("--ratio", required=True, type=float,
                   help="Fraction of train samples to include in the bank (e.g. 0.05, 0.20)")
    p.add_argument("--seed", type=int, default=666,
                   help="Random seed for bank selection. Keep identical across ratios to get nested banks.")
    p.add_argument("--device", type=str, default="cuda:0")
    p.add_argument("--model_name", type=str, default="dinov2_vits14")
    return p.parse_args()


def main():
    args = parse_args()
    build_bank(
        source_dir=args.source_dir,
        output_dir=args.output_dir,
        ratio=args.ratio,
        seed=args.seed,
        device=args.device,
        model_name=args.model_name,
    )


if __name__ == "__main__":
    main()