evaluate_model.py

import argparse
import json
from pathlib import Path

import cv2
import numpy as np
import torch
from ultralytics import YOLO


def evaluate_model(model_path, data_path, split, image_size, batch_size, device, output_path, plots):
    model = YOLO(model_path)
    metrics = model.val(
        data=data_path,
        split=split,
        imgsz=image_size,
        batch=batch_size,
        device=device,
        plots=plots,
        verbose=False,
    )

    box_metrics = metrics.box
    class_names = [model.names[index] for index in sorted(model.names)]
    report = {
        "model": str(model_path),
        "data": str(data_path),
        "split": split,
        "classes": class_names,
        "aggregate": {
            "precision": round(float(box_metrics.mp), 6),
            "recall": round(float(box_metrics.mr), 6),
            "f1": round(float(2 * box_metrics.mp * box_metrics.mr / (box_metrics.mp + box_metrics.mr)), 6),
            "map50": round(float(box_metrics.map50), 6),
            "map50_95": round(float(box_metrics.map), 6),
            "map75": round(float(box_metrics.map75), 6),
        },
        "per_class": [],
    }

    for index, class_name in enumerate(class_names):
        report["per_class"].append(
            {
                "class": class_name,
                "precision": round(float(box_metrics.p[index]), 6),
                "recall": round(float(box_metrics.r[index]), 6),
                "f1": round(float(box_metrics.f1[index]), 6),
                "map50": round(float(box_metrics.ap50[index]), 6),
                "map50_95": round(float(box_metrics.ap[index]), 6),
            }
        )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(json.dumps(report, indent=2), encoding="utf-8")

    print(f"Model: {model_path}")
    print(f"Split: {split}")
    print("\nAggregate metrics")
    for metric_name, value in report["aggregate"].items():
        print(f"{metric_name}: {value:.4f} ({value * 100:.2f}%)")

    print("\nPer-class metrics")
    print("Class | Precision | Recall | F1 | mAP50 | mAP50-95")
    for class_metrics in report["per_class"]:
        print(
            f"{class_metrics['class']} | "
            f"{class_metrics['precision']:.4f} | "
            f"{class_metrics['recall']:.4f} | "
            f"{class_metrics['f1']:.4f} | "
            f"{class_metrics['map50']:.4f} | "
            f"{class_metrics['map50_95']:.4f}"
        )

    print(f"\nJSON report saved to: {output_path}")


def _make_gradcam(model, image, image_size, device):
    """Create a Grad-CAM overlay from the last feature map before Detect."""
    network = model.model
    target_layer = network.model[-2]
    activations = []
    gradients = []

    def save_activation(_module, _inputs, output):
        activations.append(output)

        def save_gradient(gradient):
            gradients.append(gradient)

        output.register_hook(save_gradient)

    handle = target_layer.register_forward_hook(save_activation)
    try:
        with torch.enable_grad():
            resized = cv2.resize(image, (image_size, image_size), interpolation=cv2.INTER_LINEAR)
            tensor = torch.from_numpy(resized[:, :, ::-1].copy()).permute(2, 0, 1).float().unsqueeze(0) / 255.0
            tensor = tensor.to(device).requires_grad_(True)
            network.zero_grad(set_to_none=True)
            prediction = network(tensor)
            prediction_tensor = prediction[0] if isinstance(prediction, (tuple, list)) else prediction
            score = prediction_tensor.max()
            score.backward()

            activation = activations[0][0]
            gradient = gradients[0][0]
            weights = gradient.mean(dim=(1, 2), keepdim=True)
            cam = torch.relu((weights * activation).sum(dim=0)).detach().cpu().numpy()
        cam = cv2.resize(cam, (image.shape[1], image.shape[0]), interpolation=cv2.INTER_LINEAR)
        cam -= cam.min()
        maximum = cam.max()
        if maximum > 0:
            cam /= maximum
        heatmap = cv2.applyColorMap(np.uint8(cam * 255), cv2.COLORMAP_JET)
        overlay = cv2.addWeighted(image, 0.55, heatmap, 0.45, 0)
        return overlay
    finally:
        handle.remove()


def analyze_images(model_path, image_dir, output_dir, count, skip, image_size, device, confidence):
    model = YOLO(model_path)
    cam_model = YOLO(model_path)
    output_dir.mkdir(parents=True, exist_ok=True)
    originals_dir = output_dir / "original"
    results_dir = output_dir / "result"
    gradcam_dir = output_dir / "gradcam"
    for directory in (originals_dir, results_dir, gradcam_dir):
        directory.mkdir(exist_ok=True)

    image_paths = sorted(
        path for path in Path(image_dir).iterdir()
        if path.suffix.lower() in {".jpg", ".jpeg", ".png", ".bmp", ".webp"}
    )[skip:skip + count]
    if not image_paths:
        raise ValueError(f"No supported images found in {image_dir}")

    analysis = []
    cam_device = torch.device("cuda" if str(device).isdigit() and torch.cuda.is_available() else "cpu")
    cam_model.model.to(cam_device)
    cam_model.model.eval()
    for image_path in image_paths:
        image = cv2.imread(str(image_path))
        if image is None:
            continue
        prediction = model.predict(source=image, conf=confidence, imgsz=image_size, device=device, verbose=False)[0]
        result_path = results_dir / image_path.name
        cv2.imwrite(str(result_path), prediction.plot())
        cv2.imwrite(str(originals_dir / image_path.name), image)
        gradcam_path = gradcam_dir / image_path.name
        cv2.imwrite(str(gradcam_path), _make_gradcam(cam_model, image, image_size, cam_device))
        detections = []
        if prediction.boxes is not None:
            for box in prediction.boxes:
                class_id = int(box.cls.item())
                detections.append({
                    "class": prediction.names[class_id],
                    "confidence": round(float(box.conf.item()), 6),
                })
        analysis.append({
            "image": str(image_path),
            "original": str(originals_dir / image_path.name),
            "result": str(result_path),
            "gradcam": str(gradcam_path),
            "detections": detections,
        })

    report_path = output_dir / "results.json"
    report_path.write_text(json.dumps(analysis, indent=2), encoding="utf-8")
    print(f"Analyzed {len(analysis)} images")
    print(f"Results saved to: {output_dir}")
    print(f"JSON summary saved to: {report_path}")


def main():
    parser = argparse.ArgumentParser(description="Evaluate a trained Ultralytics YOLO detection model.")
    parser.add_argument("--model", default="runs/plastic_detection/weights/best.pt")
    parser.add_argument("--data", default="data.yaml")
    parser.add_argument("--split", choices=("train", "val", "test"), default="test")
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--device", default="cpu", help="Use cpu, 0, 1, or another Ultralytics device value.")
    parser.add_argument("--output", type=Path, default=Path("runs/plastic_detection/metrics_test.json"))
    parser.add_argument("--plots", action="store_true", help="Save confusion matrix and PR/F1 curve plots.")
    parser.add_argument("--gradcam", action="store_true", help="Run per-image detection and Grad-CAM analysis.")
    parser.add_argument("--image-dir", type=Path, default=Path("test/images"))
    parser.add_argument("--num-images", type=int, default=10)
    parser.add_argument("--skip-images", type=int, default=0, help="Skip this many sorted images before analysis.")
    parser.add_argument("--confidence", type=float, default=0.25)
    parser.add_argument("--gradcam-output", type=Path, default=Path("runs/plastic_detection/gradcam_10"))
    args = parser.parse_args()

    model_path = Path(args.model)
    data_path = Path(args.data)
    if not model_path.exists():
        parser.error(f"Model file not found: {model_path}")
    if not data_path.exists():
        parser.error(f"Dataset file not found: {data_path}")

    if args.gradcam:
        if not args.image_dir.exists():
            parser.error(f"Image directory not found: {args.image_dir}")
        analyze_images(
            model_path=model_path,
            image_dir=args.image_dir,
            output_dir=args.gradcam_output,
            count=args.num_images,
            skip=args.skip_images,
            image_size=args.imgsz,
            device=args.device,
            confidence=args.confidence,
        )
        return

    evaluate_model(
        model_path=model_path,
        data_path=data_path,
        split=args.split,
        image_size=args.imgsz,
        batch_size=args.batch,
        device=args.device,
        output_path=args.output,
        plots=args.plots,
    )


if __name__ == "__main__":
    main()