dfod_runtime.py

from __future__ import annotations

import copy
from typing import Any, Dict

import numpy as np

from dfod_config import normalize_config
from Baseline_Perception.Main_perception import run_perception
from Baseline_Perception.embedder import finetune_adapter_from_bundle
from Baseline_Reasoning.calibration import resolve_calibration
from Baseline_Reasoning.Main_reasoning import run_reasoning_bundle_detailed
from Baseline_Reasoning.prototype_update import apply_memory_updates, select_pseudo_updates
from Baseline_Reasoning.serialization import load_memory_bundle, save_memory


DEFAULT_RUNTIME_CONFIG: Dict[str, Any] = {
    "perception": {
        "objectness_enabled": True,
    },
    "reasoning": {
        "metrics": ["mahalanobis", "cosine", "euclidean", "manhattan"],
        "primary_metric": "mahalanobis",
        "enable_metric_fusion": True,
        "score_threshold": -5000.0,
        "margin_threshold": 0.5,
    },
    "memory": {
        "max_prototypes_per_class": 2,
        "cluster_min_supports": 6,
        "min_supports_per_mode": 3,
        "variance_alpha": 0.3,
        "variance_floor": 1e-4,
    },
    "updates": {
        "enabled": True,
        "score_threshold_min": -5000.0,
        "margin_threshold_min": 0.5,
        "update_percentile": 90,
        "margin_percentile": 75,
        "det_conf_threshold": 0.70,
        "objectness_threshold": 0.20,
        "quality_threshold": 0.20,
        "max_updates_per_image": 3,
        "max_memory_per_class": 20,
        "memory_keep_policy": "recent",
        "persist_updated_memory": True,
    },
    "adapter": {
        "enabled": True,
        "train_on_support_build": True,
        "online_enabled": True,
        "online_min_examples": 3,
        "weights_path": "adapter.pt",
    },
}


def _deep_merge(base: Dict[str, Any], override: Dict[str, Any]) -> Dict[str, Any]:
    merged = copy.deepcopy(base)
    for key, value in override.items():
        if isinstance(value, dict) and isinstance(merged.get(key), dict):
            merged[key] = _deep_merge(merged[key], value)
        else:
            merged[key] = value
    return merged


def build_runtime_config(config: Dict[str, Any] | None = None) -> Dict[str, Any]:
    merged = copy.deepcopy(DEFAULT_RUNTIME_CONFIG)
    if config:
        merged = _deep_merge(merged, config)
    return normalize_config(merged)


def _count_supports(memory_bundle: Dict[str, Any]) -> Dict[str, int]:
    classes = memory_bundle.get("classes", {})
    if not isinstance(classes, dict):
        return {}
    return {
        str(class_name): int(np.asarray(entry.get("raw_embeddings", []), dtype=np.float32).shape[0])
        for class_name, entry in sorted(classes.items())
    }


def run_dfod_inference(
    image_path: str,
    memory_path: str = "memory.npz",
    memory_bundle: Dict[str, Any] | None = None,
    config: Dict[str, Any] | None = None,
    device: str = "cpu",
    conf_thresh: float = 0.25,
    max_detections: int = 20,
    apply_updates_enabled: bool | None = None,
    updated_memory_path: str | None = None,
) -> Dict[str, Any]:
    cfg = build_runtime_config(config)
    updates_enabled = cfg["updates"].get("enabled", True) if apply_updates_enabled is None else bool(apply_updates_enabled)

    current_memory = memory_bundle if memory_bundle is not None else load_memory_bundle(memory_path)
    memory_before = _count_supports(current_memory)
    calibration = resolve_calibration(config=cfg, memory_bundle=current_memory)

    detections, embeddings = run_perception(
        image_path=image_path,
        device=device,
        conf_thresh=conf_thresh,
        max_detections=max_detections,
        config=cfg,
    )

    if embeddings.shape[0] == 0:
        accepted_updates = []
        updated_memory = current_memory
        memory_after = _count_supports(updated_memory)
        return {
            "detections": detections,
            "embeddings": embeddings,
            "results": [],
            "decisions": [],
            "scores": np.zeros((0, 0), dtype=np.float32),
            "class_names": [],
            "accepted_updates": accepted_updates,
            "memory_bundle": updated_memory,
            "memory_class_counts_before": memory_before,
            "memory_class_counts_after": memory_after,
            "memory_support_total_before": sum(memory_before.values()),
            "memory_support_total_after": sum(memory_after.values()),
        }

    results, decisions, scores, class_names = run_reasoning_bundle_detailed(
        embeddings=embeddings,
        memory_bundle=current_memory,
        config=cfg,
    )

    for idx in range(min(len(detections), len(results))):
        detections[idx]["pred_label"] = results[idx]["label"]
        detections[idx]["pred_unknown"] = results[idx]["unknown"]
        detections[idx]["pred_score"] = results[idx]["similarity"]
        detections[idx]["pred_margin"] = float(decisions[idx].get("margin", 0.0))
        detections[idx]["rejection_reason"] = decisions[idx].get("rejection_reason", "accepted")

    accepted_updates = []
    updated_memory = current_memory
    if updates_enabled:
        accepted_updates = select_pseudo_updates(
            detections=detections,
            embeddings=embeddings,
            decisions=decisions,
            config=cfg,
            calibration=calibration,
        )
        if accepted_updates:
            updated_memory = apply_memory_updates(
                memory=current_memory,
                accepted_updates=accepted_updates,
                config=cfg,
            )
            updated_memory = finetune_adapter_from_bundle(
                memory_bundle=updated_memory,
                accepted_updates=accepted_updates,
                device=device,
                config=cfg,
            )
            if cfg["updates"].get("persist_updated_memory", False):
                save_path = updated_memory_path or memory_path
                save_memory(updated_memory, save_path)

    memory_after = _count_supports(updated_memory)
    return {
        "detections": detections,
        "embeddings": embeddings,
        "results": results,
        "decisions": decisions,
        "scores": scores,
        "class_names": class_names,
        "accepted_updates": accepted_updates,
        "memory_bundle": updated_memory,
        "memory_class_counts_before": memory_before,
        "memory_class_counts_after": memory_after,
        "memory_support_total_before": sum(memory_before.values()),
        "memory_support_total_after": sum(memory_after.values()),
    }