File size: 7,410 Bytes
52874b5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 | #!/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()
|