build_memory_from_support.py

import glob
import os
from typing import Dict, List

import numpy as np

from dfod_config import normalize_config
from Baseline_Perception.Main_perception import run_perception
from Baseline_Perception.cropper import crop_box
from Baseline_Perception.embedder import (
    adapter_fingerprint,
    clear_embedder_cache,
    extract_embeddings,
    finetune_adapter_with_targets,
)
from Baseline_Perception.io import load_image, make_crop_reference
from Baseline_Reasoning.memory import combine_support_weights
from Baseline_Reasoning.serialization import save_memory


SUPPORT_DIR = "support"
OUT_PATH = "memory.npz"


DEFAULT_CONFIG = {
    "adapter": {
        "enabled": True,
        "train_on_support_build": True,
        "weights_path": "adapter.pt",
    },
    "reasoning": {
        "metrics": ["mahalanobis", "cosine"],
        "primary_metric": "mahalanobis",
    },
}


def _collect_support_examples(class_dir: str, config: Dict[str, object]) -> Dict[str, object]:
    image_paths = sorted(glob.glob(os.path.join(class_dir, "*.jpg"))) + sorted(glob.glob(os.path.join(class_dir, "*.png")))
    if len(image_paths) == 0:
        raise ValueError(f"No images found in {class_dir}")

    embeddings: List[np.ndarray] = []
    crops: List[np.ndarray] = []
    qualities: List[float] = []
    crop_refs: List[str] = []

    for image_path in image_paths:
        detections, batch_embeddings = run_perception(
            image_path=image_path,
            device="cpu",
            conf_thresh=0.25,
            max_detections=10,
            config=config,
        )
        if batch_embeddings.shape[0] == 0:
            print(f"[warn] No detections for support image: {image_path}")
            continue

        best_idx = max(range(len(detections)), key=lambda idx: float(detections[idx].get("crop_quality", detections[idx]["conf"])))
        best_det = detections[best_idx]
        embeddings.append(batch_embeddings[best_idx])
        qualities.append(float(best_det.get("crop_quality", best_det["conf"])))
        crop_refs.append(make_crop_reference(image_path=image_path, box=list(best_det["box"]), origin="support"))

        image = load_image(image_path)
        crops.append(crop_box(image=image, box=list(best_det["box"])))

    if not embeddings:
        raise ValueError(f"No usable support crops collected for {class_dir}")

    X = np.stack(embeddings, axis=0).astype(np.float32)
    X /= np.linalg.norm(X, axis=1, keepdims=True)
    return {
        "embeddings": X,
        "crops": crops,
        "qualities": np.asarray(qualities, dtype=np.float32),
        "crop_refs": crop_refs,
    }


def _initial_targets(support_data: Dict[str, Dict[str, object]]) -> tuple[List[np.ndarray], np.ndarray]:
    all_crops: List[np.ndarray] = []
    all_targets: List[np.ndarray] = []
    for _, entry in sorted(support_data.items()):
        X = np.asarray(entry["embeddings"], dtype=np.float32)
        prototype = np.mean(X, axis=0)
        prototype = prototype / max(float(np.linalg.norm(prototype)), 1e-8)
        crops = list(entry["crops"])
        all_crops.extend(crops)
        all_targets.extend([prototype.astype(np.float32, copy=False) for _ in crops])
    return all_crops, np.stack(all_targets, axis=0).astype(np.float32)


def build_memory_bundle(
    support_dir: str = SUPPORT_DIR,
    config: Dict[str, object] | None = None,
) -> Dict[str, object]:
    cfg = normalize_config(config or DEFAULT_CONFIG)
    support_data: Dict[str, Dict[str, object]] = {}

    for class_name in sorted(os.listdir(support_dir)):
        class_path = os.path.join(support_dir, class_name)
        if not os.path.isdir(class_path):
            continue
        support_data[class_name] = _collect_support_examples(class_path, cfg)

    if cfg["adapter"].get("enabled", True) and cfg["adapter"].get("train_on_support_build", True):
        all_crops, targets = _initial_targets(support_data)
        finetune_adapter_with_targets(
            crops=all_crops,
            target_embeddings=targets,
            device="cpu",
            config=cfg,
        )
        clear_embedder_cache()
        for class_name, entry in support_data.items():
            entry["embeddings"] = extract_embeddings(
                crops=list(entry["crops"]),
                device="cpu",
                config=cfg,
            )

    bundle = {
        "version": 2,
        "dim": 2048,
        "embedder_fingerprint": adapter_fingerprint(cfg),
        "calibration": {},
        "metadata": {"built_from": os.path.abspath(support_dir)},
        "classes": {},
    }

    temperature = float(cfg["support_quality"].get("weight_temperature", 0.2))
    for class_name, entry in sorted(support_data.items()):
        embeddings = np.asarray(entry["embeddings"], dtype=np.float32)
        qualities = np.asarray(entry["qualities"], dtype=np.float32)
        weights = combine_support_weights(
            embeddings=embeddings,
            quality_scores=qualities,
            temperature=temperature,
        )
        bundle["classes"][class_name] = {
            "raw_embeddings": embeddings,
            "weights": weights,
            "qualities": qualities,
            "crop_refs": list(entry["crop_refs"]),
            "mode_assignments": np.zeros(embeddings.shape[0], dtype=np.int32),
            "metadata": {"support_count": int(embeddings.shape[0])},
        }
        print(f"Class {class_name}: collected {embeddings.shape[0]} supports")

    return bundle


def main():
    bundle = build_memory_bundle()
    save_memory(bundle, OUT_PATH)
    print("Saved memory to", OUT_PATH)


if __name__ == "__main__":
    main()