Add multi-model inference handler with RDD and SurfaceAI support

- business/finishing: classification models for road assessment
- rdd: road damage detection with bounding boxes (YOLO12s RDD2022)
- surfaceai: surface type, road type, and quality classification
- Automatically selects quality model based on detected surface type

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude <noreply@anthropic.com>

.gitattributes

@@ -0,0 +1 @@
+*.pt filter=lfs diff=lfs merge=lfs -text

handler.py

@@ -0,0 +1,272 @@
+"""
+Hugging Face Inference Endpoint Custom Handler
+Handles inference for multiple YOLO models:
+- business/finishing: classification models
+- rdd: road damage detection (object detection with bounding boxes)
+- surfaceai: surface type, road type, and quality classification
+"""
+
+import base64
+import io
+from typing import Any, Dict, List
+from PIL import Image
+from ultralytics import YOLO
+
+
+class EndpointHandler:
+    # SurfaceAI class mappings
+    SURFACE_TYPES = ["asphalt", "concrete", "paving_stones", "sett", "unpaved"]
+    QUALITY_LEVELS = ["excellent", "good", "intermediate", "bad", "very_bad"]
+
+    def __init__(self, path: str = ""):
+        """
+        Initialize the handler by loading all models.
+
+        Args:
+            path: Path to the model directory (provided by HF)
+        """
+        # Classification models
+        self.models = {
+            "business": YOLO(f"{path}/models/business_best.pt"),
+            "finishing": YOLO(f"{path}/models/finishing_best.pt")
+        }
+
+        # Road Damage Detection model
+        self.rdd_model = YOLO(f"{path}/models/rdd/yolo12s_RDD2022_best.pt")
+
+        # SurfaceAI models
+        self.surfaceai_models = {
+            "surface_type": YOLO(f"{path}/models/surfaceai/surface_type_v1.pt"),
+            "road_type": YOLO(f"{path}/models/surfaceai/road_type_v1.pt"),
+            "quality": {
+                "asphalt": YOLO(f"{path}/models/surfaceai/quality/surface_quality_asphalt_v1.pt"),
+                "concrete": YOLO(f"{path}/models/surfaceai/quality/surface_quality_concrete_v1.pt"),
+                "paving_stones": YOLO(f"{path}/models/surfaceai/quality/surface_quality_paving_stones_v1.pt"),
+                "sett": YOLO(f"{path}/models/surfaceai/quality/surface_quality_sett_v1.pt"),
+                "unpaved": YOLO(f"{path}/models/surfaceai/quality/surface_quality_unpaved_v1.pt"),
+            }
+        }
+
+    def _decode_image(self, image_input: Any) -> Image.Image:
+        """
+        Decode image from various input formats.
+
+        Args:
+            image_input: Base64 string, URL, or raw bytes
+
+        Returns:
+            PIL Image object
+        """
+        if isinstance(image_input, str):
+            if image_input.startswith(("http://", "https://")):
+                import requests
+                response = requests.get(image_input, timeout=30)
+                response.raise_for_status()
+                return Image.open(io.BytesIO(response.content))
+            else:
+                # Handle base64 with or without data URI prefix
+                if "base64," in image_input:
+                    image_input = image_input.split("base64,")[1]
+                image_data = base64.b64decode(image_input)
+                return Image.open(io.BytesIO(image_data))
+        elif isinstance(image_input, bytes):
+            return Image.open(io.BytesIO(image_input))
+        else:
+            raise ValueError(f"Unsupported image input type: {type(image_input)}")
+
+    def _run_classification(self, model: YOLO, image: Image.Image) -> Dict[str, Any]:
+        """Run classification inference and return formatted results."""
+        prediction = model.predict(image, verbose=False)[0]
+        probs = prediction.probs
+        top_class_id = int(probs.top1)
+        top_class_name = prediction.names[top_class_id]
+        top_confidence = float(probs.top1conf)
+
+        all_probs = {
+            prediction.names[i]: float(probs.data[i])
+            for i in range(len(probs.data))
+        }
+
+        return {
+            "class": top_class_name,
+            "class_id": top_class_id,
+            "confidence": round(top_confidence, 4),
+            "all_probs": {k: round(v, 4) for k, v in all_probs.items()}
+        }
+
+    def _run_rdd(self, image: Image.Image, conf_threshold: float = 0.25) -> Dict[str, Any]:
+        """
+        Run Road Damage Detection and return detections with bounding boxes.
+
+        Returns:
+            {
+                "detections": [
+                    {
+                        "class": "D00",
+                        "class_id": 0,
+                        "confidence": 0.85,
+                        "bbox": [x1, y1, x2, y2]
+                    },
+                    ...
+                ],
+                "count": 2
+            }
+        """
+        prediction = self.rdd_model.predict(image, verbose=False, conf=conf_threshold)[0]
+        detections = []
+
+        if prediction.boxes is not None and len(prediction.boxes) > 0:
+            for box in prediction.boxes:
+                class_id = int(box.cls[0])
+                class_name = prediction.names[class_id]
+                confidence = float(box.conf[0])
+                bbox = box.xyxy[0].tolist()  # [x1, y1, x2, y2]
+
+                detections.append({
+                    "class": class_name,
+                    "class_id": class_id,
+                    "confidence": round(confidence, 4),
+                    "bbox": [round(coord, 2) for coord in bbox]
+                })
+
+        return {
+            "detections": detections,
+            "count": len(detections)
+        }
+
+    def _run_surfaceai(self, image: Image.Image) -> Dict[str, Any]:
+        """
+        Run SurfaceAI models for surface type, road type, and quality assessment.
+
+        Returns:
+            {
+                "surface_type": {
+                    "class": "asphalt",
+                    "confidence": 0.92,
+                    "all_probs": {...}
+                },
+                "road_type": {
+                    "class": "primary",
+                    "confidence": 0.88,
+                    "all_probs": {...}
+                },
+                "surface_quality": {
+                    "class": "good",
+                    "confidence": 0.75,
+                    "all_probs": {...},
+                    "model_used": "asphalt"
+                }
+            }
+        """
+        results = {}
+
+        # Get surface type
+        surface_result = self._run_classification(
+            self.surfaceai_models["surface_type"], image
+        )
+        results["surface_type"] = surface_result
+
+        # Get road type
+        road_result = self._run_classification(
+            self.surfaceai_models["road_type"], image
+        )
+        results["road_type"] = road_result
+
+        # Get surface quality based on detected surface type
+        detected_surface = surface_result["class"].lower()
+        if detected_surface in self.surfaceai_models["quality"]:
+            quality_model = self.surfaceai_models["quality"][detected_surface]
+            quality_result = self._run_classification(quality_model, image)
+            quality_result["model_used"] = detected_surface
+            results["surface_quality"] = quality_result
+        else:
+            # Fallback to asphalt quality model if surface type not recognized
+            quality_model = self.surfaceai_models["quality"]["asphalt"]
+            quality_result = self._run_classification(quality_model, image)
+            quality_result["model_used"] = "asphalt"
+            quality_result["note"] = f"Surface type '{detected_surface}' not recognized, using asphalt model"
+            results["surface_quality"] = quality_result
+
+        return results
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Process inference request.
+
+        Expected input format:
+        {
+            "inputs": "<base64_string or URL>",
+            "parameters": {
+                "model": "business" | "finishing" | "both" | "rdd" | "surfaceai"
+                "conf_threshold": 0.25  # optional, for RDD only
+            }
+        }
+
+        Returns for business/finishing/both:
+        [
+            {
+                "business": {"class": "...", "class_id": 0, "confidence": 0.95, "all_probs": {...}},
+                "finishing": {"class": "...", "class_id": 0, "confidence": 0.92, "all_probs": {...}}
+            }
+        ]
+
+        Returns for rdd:
+        [
+            {
+                "detections": [
+                    {"class": "D00", "class_id": 0, "confidence": 0.85, "bbox": [x1, y1, x2, y2]},
+                    ...
+                ],
+                "count": 2
+            }
+        ]
+
+        Returns for surfaceai:
+        [
+            {
+                "surface_type": {"class": "asphalt", "confidence": 0.92, "all_probs": {...}},
+                "road_type": {"class": "primary", "confidence": 0.88, "all_probs": {...}},
+                "surface_quality": {"class": "good", "confidence": 0.75, "all_probs": {...}, "model_used": "asphalt"}
+            }
+        ]
+        """
+        # Get image input
+        image_input = data.get("inputs")
+        if not image_input:
+            return [{"error": "Missing required field: inputs"}]
+
+        # Get parameters
+        parameters = data.get("parameters", {})
+        model_choice = parameters.get("model", "both")
+
+        try:
+            # Decode image
+            image = self._decode_image(image_input)
+
+            # Handle RDD model
+            if model_choice == "rdd":
+                conf_threshold = parameters.get("conf_threshold", 0.25)
+                return [self._run_rdd(image, conf_threshold)]
+
+            # Handle SurfaceAI models
+            if model_choice == "surfaceai":
+                return [self._run_surfaceai(image)]
+
+            # Handle classification models (business/finishing/both)
+            if model_choice == "both":
+                models_to_run = ["business", "finishing"]
+            elif model_choice in self.models:
+                models_to_run = [model_choice]
+            else:
+                return [{"error": f"Invalid model choice: {model_choice}. Use 'business', 'finishing', 'both', 'rdd', or 'surfaceai'"}]
+
+            # Run classification inference
+            results = {}
+            for model_name in models_to_run:
+                model = self.models[model_name]
+                results[model_name] = self._run_classification(model, image)
+
+            return [results]
+
+        except Exception as e:
+            return [{"error": str(e)}]

models/business_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3ca532122b323ba29e3c280cff785fd64a62fc6fd7016f5c2239a9e690aa0abd
+size 3197634

models/finishing_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:08c62c0a9a44899427df2b33f34f84b76d0a72f760e9790733a6384bf5a46bf2
+size 3201346

models/rdd/yolo12s_RDD2022_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccf91e3199d81ee761f4a3a012da752939c9d72e9b1eaa242e508db3163270a0
+size 18966618

models/surfaceai/quality/surface_quality_asphalt_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f3bb1d9b7884fff0dbbf736e792e7dbcf149fb7c300818601bc96f042f24ec50
+size 81618348

models/surfaceai/quality/surface_quality_concrete_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:baf27927570d268cc0bb13115b7281c05aa8892dc56d559f113a11cb8ce91402
+size 81619134

models/surfaceai/quality/surface_quality_paving_stones_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dd8f7510c2cf8b4a6be0edbb66c1cd5909a8a52a06a3f6ee1b422d965039f084
+size 81623064

models/surfaceai/quality/surface_quality_sett_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7075f1e8c0d8388d7f7d49763807701f3a124a4d2e2d6f83362c514fa59ca58d
+size 81615990

models/surfaceai/quality/surface_quality_unpaved_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29c0b73838129bdb244d8b377746d8099e2f8dee4011c13ab1e0a5a3fc9a045e
+size 81618348

models/surfaceai/road_type_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:01b7fb3e5619a8afc1d6b8bc4e382eb21e21a0b5f5a3d380eea0d67bf818a8e5
+size 81643376

models/surfaceai/surface_type_v1.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1c87e8d1649b7aec19fc6927da0315f34c865cbc6913d9e0c35c430c69e7c1de
+size 81630118

requirements.txt

@@ -0,0 +1,3 @@
+ultralytics>=8.3.0
+pillow>=10.0.0
+requests>=2.31.0