Delete raven/cache.py

raven/cache.py

@@ -1,323 +0,0 @@
-from __future__ import annotations
-
-import json
-from concurrent.futures import ThreadPoolExecutor
-from pathlib import Path
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from tqdm import tqdm
-
-from .foundation import FoundationEncoder
-from .foundation_runtime import FoundationRuntime
-from .forensic_features import load_rgb, extract_handcrafted, spectral_map, camera_residual_map
-from .specialists import load_specialists
-from .transforms import SymmetricAugment
-from .utils import configure_cpu, seed_everything, stable_unit
-
-
-def _specialist_batch_scores(spec, cam, smaps: torch.Tensor, cmaps: torch.Tensor) -> np.ndarray:
-    with torch.inference_mode():
-        ps = spec(smaps)
-        pc = cam(cmaps)
-
-        def met(pred, target):
-            err = torch.abs(pred - target)
-            flat = err.flatten(1)
-            mae = flat.mean(1)
-            mse = ((pred - target) ** 2).flatten(1).mean(1)
-            cos = 1.0 - F.cosine_similarity(pred.flatten(1), target.flatten(1), dim=1)
-            q90 = torch.quantile(flat, .90, dim=1)
-            return torch.stack([mae, mse, cos, q90], 1)
-
-        return torch.cat([met(ps, smaps), met(pc, cmaps)], 1).cpu().numpy().astype(np.float32)
-
-
-def _prepare_unique_row(item):
-
-    r, augment_copies, transform, feat_cfg, map_size, reuse_semantic = item
-    im = load_rgb(r["path"])
-    sem_tensor = transform(im)
-    aug = SymmetricAugment()
-    views = []
-    for v in range(0, max(0, int(augment_copies)) + 1):
-        if v > 0 and r.get("split") != "train":
-            break
-        view = im if v == 0 else aug(im)
-        h = extract_handcrafted(view, feat_cfg)
-        views.append({
-            "view_id": v,
-            "handcrafted": h,
-            "spectral_map": spectral_map(view, map_size),
-            "camera_map": camera_residual_map(view, map_size),
-
-            "semantic_tensor": None if (v == 0 or reuse_semantic) else transform(view),
-        })
-    return r, sem_tensor, views
-
-
-def _prepare_unique_row_safe(item):
-    try:
-        return _prepare_unique_row(item), None
-    except Exception as exc:
-        return None, f"{item[0].get('path')}: {exc}"
-
-
-def _select_augment_paths(rows: list[dict], cap_per_class: int) -> set[str]:
-
-    from collections import defaultdict
-    if cap_per_class <= 0:
-        return set()
-    chosen = set()
-    for cls in (0, 1):
-        groups = defaultdict(list)
-        for r in rows:
-            if int(r.get("label", -1)) == cls and r.get("split") == "train":
-                groups[str(r.get("source", "unknown"))].append(r)
-        if not groups:
-            continue
-        for g in groups.values():
-            g.sort(key=lambda r: stable_unit(str(r.get("path", ""))))
-        keys = sorted(groups)
-        pos = {k: 0 for k in keys}
-        count = 0
-        while count < cap_per_class:
-            added = False
-            for k in keys:
-                i = pos[k]
-                if i < len(groups[k]):
-                    chosen.add(str(groups[k][i]["path"]))
-                    pos[k] = i + 1
-                    count += 1
-                    added = True
-                    if count >= cap_per_class:
-                        break
-            if not added:
-                break
-    return chosen
-
-
-def _training_coreset(rows: list[dict], cap_per_class: int) -> list[dict]:
-
-    from collections import defaultdict
-    if cap_per_class <= 0:
-        return list(rows)
-    keep = [r for r in rows if str(r.get("split", "train")) != "train"]
-    for cls in (0, 1):
-        pool = [r for r in rows if str(r.get("split", "train")) == "train" and int(r.get("label", -1)) == cls]
-        if len(pool) <= cap_per_class:
-            keep.extend(pool)
-            continue
-        groups = defaultdict(list)
-        for r in pool:
-            parent = Path(str(r.get("path", ""))).parent.name
-            key = (str(r.get("source", "unknown")), str(r.get("light_bucket", "unknown")), parent)
-            groups[key].append(r)
-        for vals in groups.values():
-            vals.sort(key=lambda r: stable_unit(str(r.get("path", ""))))
-        keys = sorted(groups, key=lambda k: stable_unit("|".join(k)))
-        pos = {k: 0 for k in keys}
-        selected = []
-        while len(selected) < cap_per_class:
-            added = False
-            for k in keys:
-                i = pos[k]
-                if i < len(groups[k]):
-                    selected.append(groups[k][i]); pos[k] = i + 1; added = True
-                    if len(selected) >= cap_per_class:
-                        break
-            if not added:
-                break
-        keep.extend(selected)
-
-    keep.sort(key=lambda r: (str(r.get("split", "train")), int(r.get("label", 0)), str(r.get("path", ""))))
-    return keep
-
-def cache_features(rows: list[dict], cfg: dict, specialists_path: str, out_path: str, augment_copies: int = 1) -> dict:
-    hc = cfg["hardware"]
-    fc = cfg["foundation"]
-    feat_cfg = cfg["features"]
-    configure_cpu(hc.get("torch_threads", 6), hc.get("interop_threads", 1))
-    seed_everything(cfg["training"].get("seed", 1337))
-
-    input_unique = len(rows)
-    train_cap = int(fc.get("train_cap_per_class", 0))
-    if bool(fc.get("train_coreset", False)) and train_cap > 0:
-        rows = _training_coreset(rows, train_cap)
-        print(f"[cache] diversity core-set: {input_unique} -> {len(rows)} unique rows (train cap={train_cap}/class; all val/test kept)")
-
-    foundation = FoundationEncoder(
-        fc.get("backend", "dinov2_vits14"),
-        bool(fc.get("allow_fallback", True)),
-        fc.get("fallback_backend", "torchvision_convnext_tiny"),
-    )
-    foundation.eval().to("cpu")
-
-
-    p = Path(out_path)
-    p.parent.mkdir(parents=True, exist_ok=True)
-    foundation_state_path = p.with_suffix(p.suffix + ".foundation.pt")
-    torch.save(
-        {"backend": foundation.backend, "dim": int(foundation.dim), "state_dict": foundation.model.state_dict()},
-        foundation_state_path,
-    )
-
-    runtime = FoundationRuntime(fc, foundation)
-    spec, cam, sck = load_specialists(specialists_path, "cpu")
-    map_size = int(sck["map_size"])
-
-
-    bsz = max(1, int(runtime.batch_size))
-    preprocess_workers = max(0, int(hc.get("preprocess_workers", 2)))
-    reuse_sem = bool(fc.get("reuse_semantic_for_augmented_views", True))
-    augment_cap = int(fc.get("augment_max_per_class", 4000))
-    augment_paths = _select_augment_paths(rows, augment_cap) if augment_copies > 0 else set()
-    print(f"[cache] image-level augmentation selected for {len(augment_paths)} train images "
-          f"(cap={augment_cap}/class, copies={augment_copies})")
-
-    store = {k: [] for k in [
-        "semantic", "camera", "spectral", "patch", "quality", "specialist",
-        "label", "source", "path", "split", "view_id", "light_bucket", "mean_luma",
-    ]}
-
-    total_unique = len(rows)
-    executor = ThreadPoolExecutor(max_workers=preprocess_workers) if preprocess_workers > 1 else None
-
-    try:
-        for start in tqdm(range(0, total_unique, bsz), desc=f"cache-features[{runtime.backend}]"):
-            chunk_rows = rows[start:start + bsz]
-            jobs = [
-                (
-                    r,
-                    augment_copies if str(r.get("path")) in augment_paths else 0,
-                    foundation.transform, feat_cfg, map_size, reuse_sem
-                )
-                for r in chunk_rows
-            ]
-            if executor is not None:
-                results = list(executor.map(_prepare_unique_row_safe, jobs))
-                prepared = []
-                for item, err in results:
-                    if err:
-                        print(f"[warning] skip {err}")
-                    elif item is not None:
-                        prepared.append(item)
-            else:
-                prepared = []
-                for j in jobs:
-                    try:
-                        prepared.append(_prepare_unique_row(j))
-                    except Exception as exc:
-                        print(f"[warning] skip {j[0].get('path')}: {exc}")
-
-            if not prepared:
-                continue
-            if not prepared:
-                continue
-
-            base_tensors = torch.stack([x[1] for x in prepared]).float()
-            semantic_base = runtime.infer(base_tensors)
-
-
-            semantic_aug = {}
-            if not reuse_sem:
-                aug_tensors = []
-                aug_keys = []
-                for pi, (_, _, views) in enumerate(prepared):
-                    for vv in views:
-                        if vv["view_id"] > 0 and vv["semantic_tensor"] is not None:
-                            aug_tensors.append(vv["semantic_tensor"])
-                            aug_keys.append((pi, vv["view_id"]))
-                for astart in range(0, len(aug_tensors), bsz):
-                    aa = torch.stack(aug_tensors[astart:astart + bsz]).float()
-                    yy = runtime.infer(aa)
-                    for off in range(len(yy)):
-                        semantic_aug[aug_keys[astart + off]] = yy[off]
-
-
-            flat_views = []
-            smaps = []
-            cmaps = []
-            for pi, (r, _, views) in enumerate(prepared):
-                for vv in views:
-                    flat_views.append((pi, r, vv))
-                    smaps.append(torch.from_numpy(vv["spectral_map"])[None])
-                    cmaps.append(torch.from_numpy(vv["camera_map"])[None])
-            ss = _specialist_batch_scores(
-                spec,
-                cam,
-                torch.stack(smaps).float(),
-                torch.stack(cmaps).float(),
-            )
-
-            for i, (pi, r, vv) in enumerate(flat_views):
-                v = int(vv["view_id"])
-                sem = semantic_base[pi] if (v == 0 or reuse_sem) else semantic_aug[(pi, v)]
-                h = vv["handcrafted"]
-                store["semantic"].append(sem)
-                store["camera"].append(h["camera"])
-                store["spectral"].append(h["spectral"])
-                store["patch"].append(h["patch"])
-                store["quality"].append(h["quality"])
-                store["specialist"].append(ss[i])
-                store["label"].append(int(r["label"]))
-                store["source"].append(str(r.get("source", "unknown")))
-                store["path"].append(str(r["path"]))
-                store["split"].append(str(r.get("split", "train")))
-                store["view_id"].append(v)
-                store["light_bucket"].append(str(r.get("light_bucket", "unknown")))
-                store["mean_luma"].append(float(h["quality"][0]))
-    finally:
-        if executor is not None:
-            executor.shutdown(wait=True)
-        runtime.release_accelerator_memory()
-
-    arrays = {}
-    for k, v in store.items():
-        if k in {"semantic", "camera", "spectral", "patch", "quality", "specialist"}:
-            dtype = np.float16 if (k == "semantic" and bool(fc.get("cache_semantic_float16", True))) else np.float32
-            arrays[k] = np.stack(v).astype(dtype)
-        elif k in {"label", "view_id"}:
-            arrays[k] = np.asarray(v, dtype=np.int64)
-        elif k == "mean_luma":
-            arrays[k] = np.asarray(v, dtype=np.float32)
-        else:
-            arrays[k] = np.asarray(v, dtype=str)
-
-
-    if bool(fc.get("compress_feature_cache", False)):
-        np.savez_compressed(p, **arrays)
-    else:
-        np.savez(p, **arrays)
-
-    meta = {
-        "version": 5,
-        "foundation_requested": fc.get("backend"),
-        "foundation_actual": foundation.backend,
-        "foundation_runtime": runtime.backend,
-        "foundation_dim": int(foundation.dim),
-        "foundation_state": str(foundation_state_path.resolve()),
-        "specialists": str(Path(specialists_path).resolve()),
-        "augment_copies": int(augment_copies),
-        "augment_selected_images": int(len(augment_paths)),
-        "augment_max_per_class": int(augment_cap),
-        "reuse_semantic_for_augmented_views": reuse_sem,
-        "unique_images": int(total_unique),
-        "input_unique_images": int(input_unique),
-        "train_cap_per_class": int(train_cap),
-        "train_coreset": bool(fc.get("train_coreset", False)),
-        "count": int(len(arrays["label"])),
-        "dims": {k: int(arrays[k].shape[1]) for k in ["semantic", "camera", "spectral", "patch", "quality", "specialist"]},
-    }
-    p.with_suffix(p.suffix + ".meta.json").write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")
-    print(
-        f"[saved] {p} ({meta['count']} views from {meta['unique_images']} unique images, "
-        f"foundation={foundation.backend}, runtime={runtime.backend}, reuse_aug_semantic={reuse_sem})"
-    )
-    return meta
-
-
-def load_cache(path: str) -> dict[str, np.ndarray]:
-    z = np.load(path, allow_pickle=False)
-    return {k: z[k] for k in z.files}