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()