#!/usr/bin/env python3 """ Build a template bank (template/*.npz + template_summary.json) for a raw dataset, matching the schema used by existing _10percent_processed directories. Inputs: A raw dataset dir of the form: / 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: / 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. _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()