from __future__ import annotations import time from datetime import datetime, timezone from pathlib import Path from typing import Callable import torch from torch.utils.data import DataLoader, Subset from datasets.cd_dataset import CDDataset from utils.metrics import BinaryMetrics, BoundaryMetrics, normalize_binary_prediction from utils.model_adapters import BaseModelAdapter from utils.dataset_cache import dataloader_kwargs from utils.profiling import GpuProfiler, ProfilingUnavailable, count_flops, count_parameters from utils.qualitative import ( denormalize, manifest_ids, rank_for_sample, safe_sample_id, save_binary_prediction, save_probability_map, save_visual_panel, select_or_load_manifest, ) from utils.results_writer import append_to_comparison_table, save_metrics ROOT = Path(__file__).resolve().parents[1] def load_state_dict(checkpoint_path: Path) -> dict: checkpoint = torch.load(checkpoint_path, map_location="cpu") if isinstance(checkpoint, dict): for key in ("model_state_dict", "state_dict", "model"): if key in checkpoint and isinstance(checkpoint[key], dict): return checkpoint[key] if all(torch.is_tensor(v) for v in checkpoint.values()): return checkpoint raise RuntimeError(f"Checkpoint {checkpoint_path} does not contain a recognized PyTorch state_dict.") def evaluate_torch_model( *, model_name: str, dataset_cfg: dict, model: torch.nn.Module, checkpoint_path: Path, forward_fn: Callable[[torch.nn.Module, torch.Tensor, torch.Tensor], torch.Tensor], device: torch.device, batch_size: int | None = None, max_batches: int | None = None, strict_profiling: bool = True, output_dir: Path | None = None, ) -> tuple[dict, int]: dataset_name = dataset_cfg["name"] out_dir = output_dir or ROOT / "results" / model_name / dataset_name pred_dir = out_dir / "predictions" / "test" prob_dir = out_dir / "predictions" / "test_prob" visual_dir = out_dir / "visuals" / "selected_20" eval_cfg = dataset_cfg.get("eval", {}) threshold = float(eval_cfg.get("threshold", 0.5)) boundary_tolerance = int(eval_cfg.get("boundary_tolerance", 2)) state = load_state_dict(checkpoint_path) model.load_state_dict(state, strict=True) model.to(device) model.eval() ds = CDDataset(dataset_cfg["data_root"], "test", cfg=dataset_cfg, return_format="tuple") if max_batches is not None: ds_for_loader = Subset(ds, range(min(len(ds), max_batches * int(batch_size or dataset_cfg.get("batch_size", 1))))) else: ds_for_loader = ds loader = DataLoader( ds_for_loader, batch_size=int(batch_size or dataset_cfg.get("batch_size", 8)), shuffle=False, **dataloader_kwargs(dataset_cfg, torch.cuda.is_available()), ) manifest = select_or_load_manifest(dataset_cfg) selected = manifest_ids(manifest) metrics = BinaryMetrics(threshold=threshold) boundary = BoundaryMetrics(tolerance=boundary_tolerance) param_metrics = count_parameters(model) flops_metrics: dict[str, object] profiling_errors: list[str] = [] try: image_size = int(dataset_cfg.get("img_size", 256)) flops_metrics = count_flops( model, lambda: ( torch.zeros(1, 3, image_size, image_size), torch.zeros(1, 3, image_size, image_size), ), device, ) except ProfilingUnavailable as exc: flops_metrics = { "flops": None, "flops_g": None, "flops_input_shape": [[1, 3, int(dataset_cfg.get("img_size", 256)), int(dataset_cfg.get("img_size", 256))]] * 2, "flops_library": None, "flops_error": str(exc), } profiling_errors.append(str(exc)) warmup_batches = min(2, len(loader)) timed_images = 0 model_time = 0.0 end_to_end_start = time.perf_counter() mean_a = dataset_cfg.get("mean_a", [0.485, 0.456, 0.406]) std_a = dataset_cfg.get("std_a", [0.229, 0.224, 0.225]) mean_b = dataset_cfg.get("mean_b", mean_a) std_b = dataset_cfg.get("std_b", std_a) with torch.inference_mode(), GpuProfiler(device=device, required=False) as gpu_profiler: for batch_idx, (a, b, mask, names) in enumerate(loader): a = a.to(device, non_blocking=True) b = b.to(device, non_blocking=True) mask_device = mask.to(device, non_blocking=True) if device.type == "cuda": torch.cuda.synchronize(device) start = time.perf_counter() outputs = forward_fn(model, a, b) if device.type == "cuda": torch.cuda.synchronize(device) elapsed = time.perf_counter() - start pred, prob = normalize_binary_prediction(outputs.detach().cpu(), threshold=threshold) metrics.update(outputs.detach().cpu(), mask) boundary.update(pred, mask) if batch_idx >= warmup_batches: model_time += elapsed timed_images += int(a.shape[0]) for i, sample_id in enumerate(names): clean_id = safe_sample_id(str(sample_id)) pred_i = pred[i] save_binary_prediction(pred_i, pred_dir / f"{clean_id}_pred.png") prob_i = prob[i] if prob is not None else None if prob_i is not None: save_probability_map(prob_i, prob_dir / f"{clean_id}_prob.png") if str(sample_id) in selected: rank = rank_for_sample(manifest, str(sample_id)) a_vis = denormalize(a[i].detach().cpu(), mean_a, std_a) b_vis = denormalize(b[i].detach().cpu(), mean_b, std_b) save_visual_panel( a_vis, b_vis, mask[i], pred_i, visual_dir / f"{rank:02d}_{clean_id}_panel.png", prob=prob_i, ) del mask_device end_to_end_time = time.perf_counter() - end_to_end_start split_metrics = metrics.compute() split_metrics.update(boundary.compute()) split_metrics.update(param_metrics) split_metrics.update(flops_metrics) split_metrics.update(gpu_profiler.summary()) if split_metrics.get("gpu_profiling_error"): profiling_errors.append(str(split_metrics["gpu_profiling_error"])) fps_model_only = timed_images / model_time if model_time > 0 else None fps_end_to_end = len(ds_for_loader) / end_to_end_time if end_to_end_time > 0 else None status = "complete" if not (strict_profiling and profiling_errors) else "incomplete" split_metrics.update({ "model": model_name, "dataset": dataset_name, "split": "test", "checkpoint": str(checkpoint_path), "threshold": threshold, "fps": fps_model_only, "fps_model_only": fps_model_only, "fps_end_to_end": fps_end_to_end, "num_timed_images": timed_images, "warmup_batches": warmup_batches, "timing_device": str(device), "test_num_samples": len(ds_for_loader), "visual_sample_manifest": str(ROOT / "results" / "qualitative_samples" / dataset_name / "sample_manifest.json"), "prediction_dir": str(pred_dir), "visual_dir": str(visual_dir), "timestamp": datetime.now(timezone.utc).isoformat(), "status": status, "profiling_errors": profiling_errors, }) save_metrics(model_name, dataset_name, "test", split_metrics) append_to_comparison_table() return split_metrics, 0 if status == "complete" else 1 def evaluate_with_adapter( *, model_name: str, dataset_cfg: dict, model_config: dict, adapter: BaseModelAdapter, checkpoint_path: Path, device: torch.device, batch_size: int | None = None, max_batches: int | None = None, strict_profiling: bool = True, output_dir: Path | None = None, ) -> tuple[dict, int]: if not adapter.supports_inprocess_eval: raise RuntimeError(f"{model_name} does not support in-process evaluation: {adapter.notes_or_failure_reason}") dataset_name = dataset_cfg["name"] out_dir = output_dir or ROOT / "results" / model_name / dataset_name pred_dir = out_dir / "predictions" / "test" prob_dir = out_dir / "predictions" / "test_prob" visual_dir = out_dir / "visuals" / "selected_20" eval_cfg = dataset_cfg.get("eval", {}) threshold = float(eval_cfg.get("threshold", 0.5)) boundary_tolerance = int(eval_cfg.get("boundary_tolerance", 2)) model = adapter.build_model(model_config, dataset_cfg, device) adapter.load_checkpoint(model, checkpoint_path, device) model.to(device) model.eval() ds = CDDataset(dataset_cfg["data_root"], "test", cfg=dataset_cfg, return_format="tuple") if max_batches is not None: ds_for_loader = Subset(ds, range(min(len(ds), max_batches * int(batch_size or dataset_cfg.get("batch_size", 1))))) else: ds_for_loader = ds loader = DataLoader( ds_for_loader, batch_size=int(batch_size or dataset_cfg.get("batch_size", 8)), shuffle=False, **dataloader_kwargs(dataset_cfg, torch.cuda.is_available()), ) manifest = select_or_load_manifest(dataset_cfg) selected = manifest_ids(manifest) metrics = BinaryMetrics(threshold=threshold) boundary = BoundaryMetrics(tolerance=boundary_tolerance) param_metrics = count_parameters(model) profiling_errors: list[str] = [] try: if not adapter.supports_flops: raise ProfilingUnavailable(f"{model_name} adapter does not support FLOPs: {adapter.notes_or_failure_reason}") flops_metrics = count_flops(model, lambda: adapter.get_dummy_inputs(dataset_cfg, device), device) except ProfilingUnavailable as exc: flops_metrics = { "flops": None, "flops_g": None, "flops_input_shape": None, "flops_library": None, "flops_error": str(exc), } profiling_errors.append(str(exc)) warmup_batches = min(2, len(loader)) timed_images = 0 model_time = 0.0 end_to_end_start = time.perf_counter() mean_a = dataset_cfg.get("mean_a", [0.485, 0.456, 0.406]) std_a = dataset_cfg.get("std_a", [0.229, 0.224, 0.225]) mean_b = dataset_cfg.get("mean_b", mean_a) std_b = dataset_cfg.get("std_b", std_a) with torch.inference_mode(), GpuProfiler(device=device, required=False) as gpu_profiler: for batch_idx, batch in enumerate(loader): if device.type == "cuda": torch.cuda.synchronize(device) start = time.perf_counter() raw_output = adapter.forward(model, batch, device) if device.type == "cuda": torch.cuda.synchronize(device) elapsed = time.perf_counter() - start a, b, mask, names = batch normalized = adapter.normalize_output(raw_output, batch, dataset_cfg) metrics.update(normalized.metric_tensor, mask) boundary.update(normalized.binary, mask) if batch_idx >= warmup_batches: model_time += elapsed timed_images += int(a.shape[0]) for i, sample_id in enumerate(names): clean_id = safe_sample_id(str(sample_id)) pred_i = normalized.binary[i] save_binary_prediction(pred_i, pred_dir / f"{clean_id}_pred.png") prob_i = normalized.score[i] if normalized.score is not None else None if prob_i is not None: save_probability_map(prob_i, prob_dir / f"{clean_id}_prob.png") if str(sample_id) in selected: rank = rank_for_sample(manifest, str(sample_id)) a_vis = denormalize(a[i].detach().cpu(), mean_a, std_a) b_vis = denormalize(b[i].detach().cpu(), mean_b, std_b) save_visual_panel( a_vis, b_vis, mask[i], pred_i, visual_dir / f"{rank:02d}_{clean_id}_panel.png", prob=prob_i, ) end_to_end_time = time.perf_counter() - end_to_end_start split_metrics = metrics.compute() split_metrics.update(boundary.compute()) split_metrics.update(param_metrics) split_metrics.update(flops_metrics) split_metrics.update(gpu_profiler.summary()) if split_metrics.get("gpu_profiling_error"): profiling_errors.append(str(split_metrics["gpu_profiling_error"])) fps_model_only = timed_images / model_time if model_time > 0 else None fps_end_to_end = len(ds_for_loader) / end_to_end_time if end_to_end_time > 0 else None status = "complete" if not (strict_profiling and profiling_errors) else "incomplete" split_metrics.update({ "model": model_name, "dataset": dataset_name, "split": "test", "checkpoint": str(checkpoint_path), "threshold": threshold, "fps": fps_model_only, "fps_model_only": fps_model_only, "fps_end_to_end": fps_end_to_end, "num_timed_images": timed_images, "warmup_batches": warmup_batches, "timing_device": str(device), "test_num_samples": len(ds_for_loader), "visual_sample_manifest": str(ROOT / "results" / "qualitative_samples" / dataset_name / "sample_manifest.json"), "prediction_dir": str(pred_dir), "visual_dir": str(visual_dir), "timestamp": datetime.now(timezone.utc).isoformat(), "status": status, "profiling_errors": profiling_errors, "adapter": { "model_class_path": adapter.model_class_path, "input_format": adapter.input_format, "output_format": adapter.output_format, "checkpoint_format": adapter.checkpoint_format, "final_output_for_metrics": adapter.final_output_for_metrics, "notes": adapter.notes_or_failure_reason, }, }) save_metrics(model_name, dataset_name, "test", split_metrics) append_to_comparison_table() return split_metrics, 0 if status == "complete" else 1