Spaces:
Sleeping
Sleeping
| """ | |
| Evaluation script to test TrafficSentinel HSUP on the provided dataset. | |
| Runs inference on Test data folders and calculates actual Precision, Recall, and F1. | |
| Saves results to data/eval_results.json. | |
| """ | |
| import os | |
| import sys | |
| import json | |
| import time | |
| from pathlib import Path | |
| import cv2 | |
| import numpy as np | |
| import logging | |
| PROJECT_ROOT = Path(__file__).parent.parent | |
| sys.path.insert(0, str(PROJECT_ROOT)) | |
| from config.settings import SETTINGS, EntityClass, ViolationType | |
| from core.scene_conditioner import SceneConditioner | |
| from core.entity_detector import EntityDetector, Detection | |
| from core.scene_graph import SceneGraph | |
| from core.violation_engine import ViolationEngine | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s") | |
| logger = logging.getLogger(__name__) | |
| def run_dataset_evaluation(): | |
| dataset_dir = PROJECT_ROOT / "Traffic Violations Analysis Dataset" / "Test data" | |
| if not dataset_dir.exists(): | |
| logger.error(f"Dataset directory not found at {dataset_dir}") | |
| return | |
| logger.info("Initializing models for evaluation...") | |
| conditioner = SceneConditioner(SETTINGS) | |
| detector = EntityDetector(SETTINGS) | |
| violation_engine = ViolationEngine(SETTINGS) | |
| # We will evaluate helmet detection (helmet vs no_helmet) and overloading | |
| results = { | |
| "helmet": {"tp": 0, "fp": 0, "tn": 0, "fn": 0, "total": 0}, | |
| "overloading": {"tp": 0, "fp": 0, "tn": 0, "fn": 0, "total": 0} | |
| } | |
| # 1. Evaluate Helmet detection | |
| # Folder 'no_helmet': we expect to find helmet violations (TP) | |
| no_helmet_dir = dataset_dir / "no_helmet" | |
| if no_helmet_dir.exists(): | |
| files = (list(no_helmet_dir.glob("*.jpg")) + list(no_helmet_dir.glob("*.png")))[:25] | |
| logger.info(f"Evaluating {len(files)} images in 'no_helmet' folder...") | |
| for idx, f_path in enumerate(files): | |
| frame = cv2.imread(str(f_path)) | |
| if frame is None: | |
| continue | |
| # Run pipeline | |
| cond_res = conditioner.condition(frame) | |
| enhanced = cond_res.enhanced_image | |
| detections = detector.detect(enhanced, use_sahi=False) | |
| # Crop riders and detect helmets/seatbelts based on vehicle association | |
| rider_nodes = [d for d in detections if d.entity_class == EntityClass.PERSON] | |
| if rider_nodes: | |
| from core.scene_graph import compute_iou, compute_overlap_ratio, is_four_wheeler_occupant | |
| from config.settings import TWO_WHEELER_CLASSES, FOUR_WHEELER_CLASSES | |
| helmet_crops = [] | |
| seatbelt_crops = [] | |
| vehicles = [d for d in detections if d.entity_class is not None and d.entity_class in (TWO_WHEELER_CLASSES | FOUR_WHEELER_CLASSES)] | |
| for r in rider_nodes: | |
| is_two_wheeler_rider = False | |
| is_occupant = False | |
| for v in vehicles: | |
| iou = compute_iou(r.bbox, v.bbox) | |
| overlap = compute_overlap_ratio(r.bbox, v.bbox) | |
| score = max(iou, overlap) | |
| if v.entity_class in TWO_WHEELER_CLASSES: | |
| if score >= 0.15: | |
| is_two_wheeler_rider = True | |
| elif v.entity_class in FOUR_WHEELER_CLASSES: | |
| if is_four_wheeler_occupant(r.bbox, v.bbox): | |
| is_occupant = True | |
| if is_two_wheeler_rider: | |
| helmet_crops.append(r.bbox) | |
| if is_occupant: | |
| seatbelt_crops.append(r.bbox) | |
| if helmet_crops: | |
| helmet_detections = detector.detect_helmets(enhanced, helmet_crops) | |
| detections.extend(helmet_detections) | |
| if seatbelt_crops: | |
| seatbelt_detections = detector.detect_seatbelts(enhanced, seatbelt_crops) | |
| detections.extend(seatbelt_detections) | |
| sg = SceneGraph(SETTINGS) | |
| sg.build(detections, frame_shape=enhanced.shape[:2]) | |
| violations = violation_engine.analyze(sg, enhanced) | |
| has_helmet_violation = any(v.violation_type == ViolationType.HELMET_NON_COMPLIANCE for v in violations) | |
| results["helmet"]["total"] += 1 | |
| if has_helmet_violation: | |
| results["helmet"]["tp"] += 1 | |
| else: | |
| results["helmet"]["fn"] += 1 | |
| if (idx + 1) % 20 == 0: | |
| logger.info(f"Processed {idx + 1}/{len(files)} no_helmet images") | |
| # Folder 'helmet': we expect NO helmet violations (TN) | |
| helmet_dir = dataset_dir / "helmet" | |
| if helmet_dir.exists(): | |
| files = (list(helmet_dir.glob("*.jpg")) + list(helmet_dir.glob("*.png")))[:25] | |
| logger.info(f"Evaluating {len(files)} images in 'helmet' folder...") | |
| for idx, f_path in enumerate(files): | |
| frame = cv2.imread(str(f_path)) | |
| if frame is None: | |
| continue | |
| # Run pipeline | |
| cond_res = conditioner.condition(frame) | |
| enhanced = cond_res.enhanced_image | |
| detections = detector.detect(enhanced, use_sahi=False) | |
| # Crop riders and detect helmets/seatbelts based on vehicle association | |
| rider_nodes = [d for d in detections if d.entity_class == EntityClass.PERSON] | |
| if rider_nodes: | |
| from core.scene_graph import compute_iou, compute_overlap_ratio, is_four_wheeler_occupant | |
| from config.settings import TWO_WHEELER_CLASSES, FOUR_WHEELER_CLASSES | |
| helmet_crops = [] | |
| seatbelt_crops = [] | |
| vehicles = [d for d in detections if d.entity_class is not None and d.entity_class in (TWO_WHEELER_CLASSES | FOUR_WHEELER_CLASSES)] | |
| for r in rider_nodes: | |
| is_two_wheeler_rider = False | |
| is_occupant = False | |
| for v in vehicles: | |
| iou = compute_iou(r.bbox, v.bbox) | |
| overlap = compute_overlap_ratio(r.bbox, v.bbox) | |
| score = max(iou, overlap) | |
| if v.entity_class in TWO_WHEELER_CLASSES: | |
| if score >= 0.15: | |
| is_two_wheeler_rider = True | |
| elif v.entity_class in FOUR_WHEELER_CLASSES: | |
| if is_four_wheeler_occupant(r.bbox, v.bbox): | |
| is_occupant = True | |
| if is_two_wheeler_rider: | |
| helmet_crops.append(r.bbox) | |
| if is_occupant: | |
| seatbelt_crops.append(r.bbox) | |
| if helmet_crops: | |
| helmet_detections = detector.detect_helmets(enhanced, helmet_crops) | |
| detections.extend(helmet_detections) | |
| if seatbelt_crops: | |
| seatbelt_detections = detector.detect_seatbelts(enhanced, seatbelt_crops) | |
| detections.extend(seatbelt_detections) | |
| sg = SceneGraph(SETTINGS) | |
| sg.build(detections, frame_shape=enhanced.shape[:2]) | |
| violations = violation_engine.analyze(sg, enhanced) | |
| has_helmet_violation = any(v.violation_type == ViolationType.HELMET_NON_COMPLIANCE for v in violations) | |
| results["helmet"]["total"] += 1 | |
| if has_helmet_violation: | |
| results["helmet"]["fp"] += 1 | |
| else: | |
| results["helmet"]["tn"] += 1 | |
| if (idx + 1) % 20 == 0: | |
| logger.info(f"Processed {idx + 1}/{len(files)} helmet images") | |
| # 2. Evaluate Overloading detection | |
| # Folder 'overloading': we expect overloading/triple riding violations (TP) | |
| overloading_dir = dataset_dir / "overloading" | |
| if overloading_dir.exists(): | |
| files = (list(overloading_dir.glob("*.jpg")) + list(overloading_dir.glob("*.png")))[:25] | |
| logger.info(f"Evaluating {len(files)} images in 'overloading' folder...") | |
| for idx, f_path in enumerate(files): | |
| frame = cv2.imread(str(f_path)) | |
| if frame is None: | |
| continue | |
| # Run pipeline | |
| cond_res = conditioner.condition(frame) | |
| enhanced = cond_res.enhanced_image | |
| detections = detector.detect(enhanced, use_sahi=False) | |
| # Crop riders and detect helmets/seatbelts based on vehicle association | |
| rider_nodes = [d for d in detections if d.entity_class == EntityClass.PERSON] | |
| if rider_nodes: | |
| from core.scene_graph import compute_iou, compute_overlap_ratio, is_four_wheeler_occupant | |
| from config.settings import TWO_WHEELER_CLASSES, FOUR_WHEELER_CLASSES | |
| helmet_crops = [] | |
| seatbelt_crops = [] | |
| vehicles = [d for d in detections if d.entity_class is not None and d.entity_class in (TWO_WHEELER_CLASSES | FOUR_WHEELER_CLASSES)] | |
| for r in rider_nodes: | |
| is_two_wheeler_rider = False | |
| is_occupant = False | |
| for v in vehicles: | |
| iou = compute_iou(r.bbox, v.bbox) | |
| overlap = compute_overlap_ratio(r.bbox, v.bbox) | |
| score = max(iou, overlap) | |
| if v.entity_class in TWO_WHEELER_CLASSES: | |
| if score >= 0.15: | |
| is_two_wheeler_rider = True | |
| elif v.entity_class in FOUR_WHEELER_CLASSES: | |
| if is_four_wheeler_occupant(r.bbox, v.bbox): | |
| is_occupant = True | |
| if is_two_wheeler_rider: | |
| helmet_crops.append(r.bbox) | |
| if is_occupant: | |
| seatbelt_crops.append(r.bbox) | |
| if helmet_crops: | |
| helmet_detections = detector.detect_helmets(enhanced, helmet_crops) | |
| detections.extend(helmet_detections) | |
| if seatbelt_crops: | |
| seatbelt_detections = detector.detect_seatbelts(enhanced, seatbelt_crops) | |
| detections.extend(seatbelt_detections) | |
| sg = SceneGraph(SETTINGS) | |
| sg.build(detections, frame_shape=enhanced.shape[:2]) | |
| violations = violation_engine.analyze(sg, enhanced) | |
| has_overloading = any(v.violation_type == ViolationType.TRIPLE_RIDING for v in violations) | |
| results["overloading"]["total"] += 1 | |
| if has_overloading: | |
| results["overloading"]["tp"] += 1 | |
| else: | |
| results["overloading"]["fn"] += 1 | |
| if (idx + 1) % 20 == 0: | |
| logger.info(f"Processed {idx + 1}/{len(files)} overloading images") | |
| # For overloading FP/TN, we use the 'helmet' and 'no_helmet' folders as negative cases (normal capacity) | |
| # We check how many normal images got flagged as triple riding (FP) vs ignored (TN) | |
| for idx, f_path in enumerate(list(helmet_dir.glob("*.jpg"))[:12] + list(no_helmet_dir.glob("*.jpg"))[:13]): | |
| frame = cv2.imread(str(f_path)) | |
| if frame is None: | |
| continue | |
| cond_res = conditioner.condition(frame) | |
| enhanced = cond_res.enhanced_image | |
| detections = detector.detect(enhanced, use_sahi=False) | |
| sg = SceneGraph(SETTINGS) | |
| sg.build(detections) | |
| violations = violation_engine.analyze(sg, enhanced) | |
| has_overloading = any(v.violation_type == ViolationType.TRIPLE_RIDING for v in violations) | |
| if has_overloading: | |
| results["overloading"]["fp"] += 1 | |
| else: | |
| results["overloading"]["tn"] += 1 | |
| # Calculate final metrics | |
| metrics = {} | |
| for task_name, counts in results.items(): | |
| tp, fp, tn, fn = counts["tp"], counts["fp"], counts["tn"], counts["fn"] | |
| precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0 | |
| recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0 | |
| f1 = 2 * (precision * recall) / (precision + recall) if (precision + recall) > 0 else 0.0 | |
| accuracy = (tp + tn) / (tp + fp + tn + fn) if (tp + fp + tn + fn) > 0 else 0.0 | |
| metrics[task_name] = { | |
| "tp": tp, "fp": fp, "tn": tn, "fn": fn, | |
| "precision": float(precision), | |
| "recall": float(recall), | |
| "f1": float(f1), | |
| "accuracy": float(accuracy) | |
| } | |
| logger.info(f"{task_name.upper()} Metrics: Precision={precision:.2%}, Recall={recall:.2%}, F1={f1:.2%}, Accuracy={accuracy:.2%}") | |
| # Save to JSON | |
| output_path = PROJECT_ROOT / "data" / "eval_results.json" | |
| output_path.parent.mkdir(parents=True, exist_ok=True) | |
| with open(output_path, "w") as f: | |
| json.dump(metrics, f, indent=4) | |
| logger.info(f"Saved evaluation results to {output_path}") | |
| if __name__ == "__main__": | |
| run_dataset_evaluation() | |