Spaces:
Running on Zero
Running on Zero
| """Inference interface for the accident detector. | |
| Defines the Phase 5 output contract. Detector confidence is returned as | |
| `accident_detection_confidence` and is NOT claimed to be a calibrated probability. | |
| """ | |
| from __future__ import annotations | |
| from pathlib import Path | |
| from typing import Any | |
| import numpy as np | |
| import torch | |
| from PIL import Image | |
| from . import CLASS_NAMES | |
| from .labels import MODEL_LABEL_BACKGROUND, model_to_dataset_labels | |
| from .model import build_detector | |
| from .transforms import ResizeTransform, clamp_boxes | |
| from .utils import get_device | |
| class AccidentDetector: | |
| """Callable detector wrapper that implements the Phase 5 inference contract. | |
| The checkpoint stores a 3-class model head (background + accident + | |
| non_accident in model label space); predictions are mapped back to dataset | |
| semantics (0 = accident, 1 = non_accident) before output. | |
| """ | |
| def __init__( | |
| self, | |
| checkpoint_path: str | Path, | |
| num_classes: int | None = None, | |
| image_size: tuple[int, int] = (320, 320), | |
| confidence_threshold: float = 0.5, | |
| device: torch.device | None = None, | |
| ): | |
| self.image_size = image_size | |
| self.confidence_threshold = confidence_threshold | |
| self.device = device if device is not None else get_device() | |
| self.resize = ResizeTransform(target_size=image_size) | |
| state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True) | |
| # The 2-class Phase 5 baseline and the 3-class Phase 5B checkpoints | |
| # differ only in the box head width; infer it from the checkpoint. | |
| inferred = state_dict["roi_heads.box_predictor.cls_score.bias"].shape[0] | |
| self.num_classes = num_classes if num_classes is not None else inferred | |
| self.model = build_detector(num_classes=self.num_classes, pretrained=False) | |
| self.model.load_state_dict(state_dict) | |
| self.model.to(self.device) | |
| self.model.eval() | |
| def predict( | |
| self, | |
| image: Image.Image | np.ndarray | torch.Tensor, | |
| score_threshold: float | None = None, | |
| ) -> dict[str, Any]: | |
| """Run inference on a single image and return the contract output. | |
| Output fields: | |
| - detector_version: static version string | |
| - image_size: (width, height) | |
| - detections: list of {bbox, category, confidence} | |
| - accident_detection_confidence: aggregated detector confidence | |
| """ | |
| if isinstance(image, np.ndarray): | |
| image = Image.fromarray(image).convert("RGB") | |
| elif isinstance(image, torch.Tensor): | |
| image = Image.fromarray( | |
| (image.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) | |
| ).convert("RGB") | |
| orig_size = image.size | |
| tensor, _ = self.resize(image, np.zeros((0, 4), dtype=np.float32)) | |
| tensor = tensor.to(self.device) | |
| outputs = self.model([tensor]) | |
| output = outputs[0] | |
| thr = score_threshold if score_threshold is not None else self.confidence_threshold | |
| keep = (output["scores"] >= thr) & (output["labels"] != MODEL_LABEL_BACKGROUND) | |
| boxes_xyxy = output["boxes"][keep].cpu().numpy() | |
| # Model labels (1/2) -> dataset semantics (0/1) for category naming. | |
| labels = model_to_dataset_labels(output["labels"][keep].cpu().numpy()) | |
| scores = output["scores"][keep].cpu().numpy() | |
| # Clamp boxes to image bounds for safety. | |
| boxes_xyxy = clamp_boxes(boxes_xyxy, self.image_size[0], self.image_size[1]) | |
| detections = [ | |
| { | |
| "bbox": [round(float(v), 2) for v in box], | |
| "category": CLASS_NAMES.get(int(label), str(label)), | |
| "confidence": round(float(score), 4), | |
| } | |
| for box, label, score in zip(boxes_xyxy, labels, scores) | |
| ] | |
| accident_scores = [ | |
| d["confidence"] for d in detections if d["category"] == "accident" | |
| ] | |
| accident_detection_confidence = float(max(accident_scores)) if accident_scores else 0.0 | |
| return { | |
| "detector_version": "accident_detector_v1", | |
| "image_size": [orig_size[0], orig_size[1]], | |
| "detections": detections, | |
| "accident_detection_confidence": accident_detection_confidence, | |
| } | |
| def __call__(self, *args, **kwargs) -> dict[str, Any]: | |
| return self.predict(*args, **kwargs) | |
| def infer_probability( | |
| image: Image.Image | np.ndarray | torch.Tensor, | |
| checkpoint_path: str | Path, | |
| **kwargs: Any, | |
| ) -> dict[str, Any]: | |
| """Convenience helper to run a single-image prediction.""" | |
| detector = AccidentDetector(checkpoint_path, **kwargs) | |
| return detector.predict(image) | |