handler.py

"""
HuggingFace Inference Endpoint handler for SurfaceAI models.

This handler loads all 7 SurfaceAI models and performs hierarchical classification:
1. Road type classification
2. Surface type classification
3. Surface quality regression (model selected based on surface type)

Deploy by creating an Inference Endpoint pointing to this repo.
"""

import base64
import io
import logging
from pathlib import Path
from typing import Any, Dict, List

import torch
from huggingface_hub import hf_hub_download
from PIL import Image
from torchvision import models, transforms
from torch import nn, Tensor

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants from original SurfaceAI
NORM_MEAN = [0.42834484577178955, 0.4461250305175781, 0.4350937306880951]
NORM_SD = [0.22991590201854706, 0.23555299639701843, 0.26348039507865906]
CROP_LOWER_MIDDLE_HALF = "lower_middle_half"
CROP_LOWER_HALF = "lower_half"

# Model configuration
MODEL_CONFIG = {
    "hf_repo": "SurfaceAI/models-moved",
    "models": {
        "road_type": "v1/road_type_v1.pt",
        "surface_type": "v1/surface_type_v1.pt",
        "surface_quality": {
            "asphalt": "v1/surface_quality_asphalt_v1.pt",
            "concrete": "v1/surface_quality_concrete_v1.pt",
            "paving_stones": "v1/surface_quality_paving_stones_v1.pt",
            "sett": "v1/surface_quality_sett_v1.pt",
            "unpaved": "v1/surface_quality_unpaved_v1.pt",
        }
    },
    "transform_surface": {
        "resize": 256,
        "crop": CROP_LOWER_MIDDLE_HALF,
        "normalize": (NORM_MEAN, NORM_SD),
    },
    "transform_road_type": {
        "resize": 256,
        "crop": CROP_LOWER_HALF,
        "normalize": (NORM_MEAN, NORM_SD),
    },
}

# Quality class mapping
QUALITY_CLASSES = {
    1: "excellent",
    2: "good",
    3: "intermediate",
    4: "bad",
    5: "very_bad",
}


class CustomEfficientNetV2SLinear(nn.Module):
    """EfficientNetV2-S with linear classifier for classification/regression."""

    def __init__(self, num_classes, avg_pool=1):
        super().__init__()
        model = models.efficientnet_v2_s(weights="IMAGENET1K_V1")
        in_features = model.classifier[-1].in_features * (avg_pool * avg_pool)
        fc = nn.Linear(in_features, num_classes, bias=True)
        model.classifier[-1] = fc

        self.features = model.features
        self.avgpool = nn.AdaptiveAvgPool2d(avg_pool)
        self.classifier = model.classifier
        self.is_regression = num_classes == 1

    def forward(self, x: Tensor) -> Tensor:
        x = self.features(x)
        x = self.avgpool(x)
        x = torch.flatten(x, 1)
        x = self.classifier(x)
        return x

    def get_class_probabilities(self, x):
        if self.is_regression:
            return x.flatten()
        return nn.functional.softmax(x, dim=1)


class EndpointHandler:
    """HuggingFace Inference Endpoint handler for SurfaceAI."""

    def __init__(self, path: str = ""):
        """
        Initialize handler and load all models.

        Args:
            path: Path to model directory (provided by HF Inference Endpoints)
        """
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"Using device: {self.device}")

        self.models = {}
        self.class_mappings = {}
        self._load_all_models()

        # Pre-build transforms
        self.transform_surface = self._build_transform(MODEL_CONFIG["transform_surface"])
        self.transform_road_type = self._build_transform(MODEL_CONFIG["transform_road_type"])

    def _download_model(self, filename: str) -> str:
        """Download model from HuggingFace Hub."""
        return hf_hub_download(
            repo_id=MODEL_CONFIG["hf_repo"],
            filename=filename,
        )

    def _load_model(self, model_path: str) -> tuple:
        """Load a single model and return (model, class_to_idx, is_regression)."""
        state = torch.load(model_path, map_location=self.device, weights_only=False)

        is_regression = state["is_regression"]
        class_to_idx = state["class_to_idx"]
        num_classes = 1 if is_regression else len(class_to_idx)

        model = CustomEfficientNetV2SLinear(num_classes=num_classes)
        model.load_state_dict(state["model_state_dict"])
        model.to(self.device)
        model.eval()

        return model, class_to_idx, is_regression

    def _load_all_models(self):
        """Load all 7 SurfaceAI models."""
        logger.info("Loading SurfaceAI models...")

        # Load road type model
        path = self._download_model(MODEL_CONFIG["models"]["road_type"])
        self.models["road_type"], self.class_mappings["road_type"], _ = self._load_model(path)
        logger.info("Loaded road_type model")

        # Load surface type model
        path = self._download_model(MODEL_CONFIG["models"]["surface_type"])
        self.models["surface_type"], self.class_mappings["surface_type"], _ = self._load_model(path)
        logger.info("Loaded surface_type model")

        # Load quality models for each surface type
        self.models["quality"] = {}
        self.class_mappings["quality"] = {}
        for surface_type, model_file in MODEL_CONFIG["models"]["surface_quality"].items():
            path = self._download_model(model_file)
            model, class_to_idx, _ = self._load_model(path)
            self.models["quality"][surface_type] = model
            self.class_mappings["quality"][surface_type] = class_to_idx
            logger.info(f"Loaded quality model for {surface_type}")

        logger.info("All models loaded successfully")

    @staticmethod
    def _custom_crop(img: Image.Image, crop_style: str) -> Image.Image:
        """Crop image according to style."""
        im_width, im_height = img.size

        if crop_style == CROP_LOWER_MIDDLE_HALF:
            top = im_height // 2
            left = im_width // 4
            height = im_height // 2
            width = im_width // 2
        elif crop_style == CROP_LOWER_HALF:
            top = im_height // 2
            left = 0
            height = im_height // 2
            width = im_width
        else:
            return img

        return img.crop((left, top, left + width, top + height))

    def _build_transform(self, config: dict) -> transforms.Compose:
        """Build torchvision transform from config."""
        transform_list = []

        if config.get("crop"):
            transform_list.append(
                transforms.Lambda(lambda img: self._custom_crop(img, config["crop"]))
            )

        if config.get("resize"):
            size = config["resize"]
            if isinstance(size, int):
                size = (size, size)
            transform_list.append(transforms.Resize(size))

        transform_list.append(transforms.ToTensor())

        if config.get("normalize"):
            transform_list.append(transforms.Normalize(*config["normalize"]))

        return transforms.Compose(transform_list)

    def _predict(self, model, data: torch.Tensor, class_to_idx: dict) -> tuple:
        """Run prediction and convert to class/value."""
        with torch.no_grad():
            outputs = model(data)
            values = model.get_class_probabilities(outputs)

        idx_to_class = {i: cls for cls, i in class_to_idx.items()}

        if len(values.shape) < 2:
            # Regression output
            classes = [
                idx_to_class[
                    min(max(int(v.round().item()), min(class_to_idx.values())),
                        max(class_to_idx.values()))
                ]
                for v in values
            ]
            values_list = values.tolist()
        else:
            # Classification output
            classes = [idx_to_class[idx.item()] for idx in torch.argmax(values, dim=1)]
            values_list = values.tolist()

        return classes, values_list

    def _process_image(self, image: Image.Image) -> dict:
        """Process a single image through all models."""
        # Ensure RGB
        if image.mode != "RGB":
            image = image.convert("RGB")

        # Road type prediction
        road_data = self.transform_road_type(image).unsqueeze(0).to(self.device)
        road_classes, road_values = self._predict(
            self.models["road_type"],
            road_data,
            self.class_mappings["road_type"]
        )

        # Surface type prediction
        surface_data = self.transform_surface(image).unsqueeze(0).to(self.device)
        surface_classes, surface_values = self._predict(
            self.models["surface_type"],
            surface_data,
            self.class_mappings["surface_type"]
        )

        # Quality prediction based on detected surface type
        surface_type = surface_classes[0]
        quality_class = None
        quality_value = None

        if surface_type in self.models["quality"]:
            quality_classes, quality_values = self._predict(
                self.models["quality"][surface_type],
                surface_data,
                self.class_mappings["quality"][surface_type]
            )
            quality_class = quality_classes[0]
            quality_value = quality_values[0]

        return {
            "road_type": road_classes[0],
            "road_type_confidence": max(road_values[0]) if isinstance(road_values[0], list) else road_values[0],
            "surface_type": surface_type,
            "surface_type_confidence": max(surface_values[0]) if isinstance(surface_values[0], list) else surface_values[0],
            "quality_class": quality_class,
            "quality_value": quality_value,
        }

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Process inference request.

        Args:
            data: Request data containing either:
                - "inputs": base64-encoded image or URL
                - "image": PIL Image (when called directly)

        Returns:
            List of prediction results
        """
        inputs = data.get("inputs", data.get("image"))

        if inputs is None:
            return [{"error": "No input provided. Send 'inputs' with base64 image or URL."}]

        try:
            # Handle different input types
            if isinstance(inputs, str):
                if inputs.startswith("data:image"):
                    # Base64 data URL
                    inputs = inputs.split(",")[1]
                    image_bytes = base64.b64decode(inputs)
                    image = Image.open(io.BytesIO(image_bytes))
                elif inputs.startswith("http"):
                    # URL - fetch it
                    import requests
                    response = requests.get(inputs, timeout=10)
                    image = Image.open(io.BytesIO(response.content))
                else:
                    # Assume raw base64
                    image_bytes = base64.b64decode(inputs)
                    image = Image.open(io.BytesIO(image_bytes))
            elif isinstance(inputs, Image.Image):
                image = inputs
            elif isinstance(inputs, bytes):
                image = Image.open(io.BytesIO(inputs))
            else:
                return [{"error": f"Unsupported input type: {type(inputs)}"}]

            result = self._process_image(image)
            return [result]

        except Exception as e:
            logger.exception("Error processing request")
            return [{"error": str(e)}]