Droid210/FleetVision

1#!/usr/bin/env python2"""Validate Model B checkpoint on validation data with PASS/FAIL thresholds."""3 4from __future__ import annotations5 6import argparse7import os8import sys9from pathlib import Path10 11import torch12from sklearn.metrics import confusion_matrix13 14 15def parse_args() -> argparse.Namespace:16    root = Path(__file__).resolve().parent17    parser = argparse.ArgumentParser(description="Test Model B weights on validation data")18    parser.add_argument(19        "--data-dir",20        type=Path,21        default=root / "data" / "model b" / "damage_assessment",22        help="Dataset root used by Model B loader",23    )24    parser.add_argument(25        "--model-path",26        type=Path,27        default=root / "weights" / "model b" / "best_damage_detector.pth",28        help="Checkpoint path",29    )30    parser.add_argument("--batch-size", type=int, default=32)31    parser.add_argument("--num-workers", type=int, default=0)32    parser.add_argument(33        "--max-batches",34        type=int,35        default=0,36        help="If > 0, evaluate only this many validation batches for a fast smoke check",37    )38    parser.add_argument("--min-acc", type=float, default=0.90)39    parser.add_argument("--min-damage-recall", type=float, default=0.85)40    parser.add_argument(41        "--device",42        type=str,43        default="auto",44        choices=["auto", "cpu", "cuda"],45        help="Device for evaluation",46    )47    return parser.parse_args()48 49 50def resolve_device(choice: str) -> torch.device:51    if choice == "cpu":52        return torch.device("cpu")53    if choice == "cuda":54        return torch.device("cuda")55 56    if torch.cuda.is_available():57        try:58            _ = torch.randn(1, device="cuda")59            return torch.device("cuda")60        except RuntimeError:61            print("WARNING: CUDA detected but unusable. Falling back to CPU.")62            return torch.device("cpu")63    return torch.device("cpu")64 65 66def main() -> int:67    args = parse_args()68 69    root = Path(__file__).resolve().parent70    if str(root) not in sys.path:71        sys.path.insert(0, str(root))72 73    from models.model_b.data import build_dataloaders74    from models.model_b.inference import load_trained_model75 76    if not args.data_dir.exists():77        print(f"FAIL: Data directory not found: {args.data_dir}")78        return 179 80    if not args.model_path.exists():81        print(f"FAIL: Checkpoint not found: {args.model_path}")82        return 183 84    num_workers = args.num_workers85    if os.name == "nt" and num_workers > 0:86        print("WARNING: On Windows, forcing num_workers=0 for stability.")87        num_workers = 088 89    train_loader, val_loader, _ = build_dataloaders(90        data_dir=args.data_dir,91        batch_size=args.batch_size,92        num_workers=num_workers,93    )94 95    device = resolve_device(args.device)96    print(f"Device: {device}")97    print(f"Validation samples: {len(val_loader.dataset)}")98 99    try:100        model, _ = load_trained_model(args.model_path, device)101    except RuntimeError as exc:102        if "CUDA" in str(exc) and device.type == "cuda":103            print("WARNING: CUDA load failed; retrying on CPU.")104            device = torch.device("cpu")105            model, _ = load_trained_model(args.model_path, device)106        else:107            raise108 109    model.eval()110 111    y_true: list[int] = []112    y_pred: list[int] = []113 114    with torch.no_grad():115        for batch_idx, (pixel_values, labels) in enumerate(val_loader):116            if args.max_batches > 0 and batch_idx >= args.max_batches:117                break118            pixel_values = pixel_values.to(device)119            labels = labels.to(device)120            logits = model(pixel_values).logits121            preds = logits.argmax(dim=1)122 123            y_true.extend(labels.cpu().tolist())124            y_pred.extend(preds.cpu().tolist())125 126    if not y_true:127        print("FAIL: No validation samples were evaluated.")128        return 1129 130    cm = confusion_matrix(y_true, y_pred, labels=[0, 1])131    tn, fp, fn, tp = cm.ravel()132 133    total = len(y_true)134    accuracy = (tp + tn) / max(1, total)135    damage_recall = tp / max(1, (tp + fn))136    damage_precision = tp / max(1, (tp + fp))137 138    print("\n=== Validation Summary ===")139    print(f"Accuracy       : {accuracy:.4f}")140    print(f"Damage Recall  : {damage_recall:.4f}")141    print(f"Damage Precision: {damage_precision:.4f}")142    print("Confusion Matrix [[TN, FP], [FN, TP]]:")143    print(cm)144 145    acc_ok = accuracy >= args.min_acc146    rec_ok = damage_recall >= args.min_damage_recall147 148    print("\n=== Threshold Check ===")149    print(f"Accuracy >= {args.min_acc:.2f}: {'PASS' if acc_ok else 'FAIL'}")150    print(f"Damage Recall >= {args.min_damage_recall:.2f}: {'PASS' if rec_ok else 'FAIL'}")151 152    if acc_ok and rec_ok:153        print("\nPASS: Model weights meet requested quality thresholds.")154        return 0155 156    print("\nFAIL: Model weights do not meet one or more thresholds.")157    return 2158 159 160if __name__ == "__main__":161    raise SystemExit(main())162