rezzzq

Claude

Fix SurfaceAI models to use EfficientNetV2 architecture

bb548d7 11 months ago

13.9 kB

"""
Hugging Face Inference Endpoint Custom Handler
Handles inference for multiple models:
- business/finishing: YOLO classification models
- rdd: YOLO road damage detection (object detection with bounding boxes)
- surfaceai: EfficientNetV2 models for surface type, road type, and quality classification
"""

import base64
import io
from typing import Any, Dict, List, Tuple
from PIL import Image
from ultralytics import YOLO
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.models import efficientnet_v2_s


class EfficientNetClassifier:
    """Wrapper for EfficientNetV2 classification models."""

    def __init__(self, model_path: str, device: str = "cpu"):
        checkpoint = torch.load(model_path, map_location=device, weights_only=False)
        self.class_to_idx = checkpoint["class_to_idx"]
        self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}
        self.num_classes = len(self.class_to_idx)
        self.is_regression = checkpoint.get("is_regression", False)
        self.device = device

        # Determine output size
        output_size = 1 if self.is_regression else self.num_classes

        # Build model
        self.model = efficientnet_v2_s(weights=None)
        self.model.classifier = nn.Sequential(
            nn.Dropout(p=0.2, inplace=True),
            nn.Linear(self.model.classifier[1].in_features, output_size)
        )
        self.model.load_state_dict(checkpoint["model_state_dict"])
        self.model.to(device)
        self.model.eval()

        # Image transforms (EfficientNetV2-S uses 384x384)
        self.transform = transforms.Compose([
            transforms.Resize((384, 384)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def predict(self, image: Image.Image) -> Tuple[str, int, float, Dict[str, float]]:
        """Run inference and return class, id, confidence, and all probabilities."""
        image = image.convert("RGB")
        input_tensor = self.transform(image).unsqueeze(0).to(self.device)

        with torch.no_grad():
            outputs = self.model(input_tensor)

        if self.is_regression:
            # Regression model: output is a quality score
            raw_score = float(outputs[0, 0])
            # Clamp to valid range based on this model's class indices
            min_idx = min(self.idx_to_class.keys())
            max_idx = max(self.idx_to_class.keys())
            score = max(min_idx, min(max_idx, raw_score))
            class_id = int(round(score))
            # Ensure class_id is valid
            if class_id not in self.idx_to_class:
                class_id = min(self.idx_to_class.keys(), key=lambda x: abs(x - score))
            class_name = self.idx_to_class[class_id]

            # Create pseudo-probabilities based on distance from score
            all_probs = {}
            for idx, name in self.idx_to_class.items():
                distance = abs(idx - score)
                all_probs[name] = max(0, 1 - distance * 0.25)

            return class_name, class_id, score, all_probs
        else:
            # Classification model
            probs = F.softmax(outputs, dim=1)[0]
            top_prob, top_idx = torch.max(probs, 0)
            top_class_id = int(top_idx)
            top_class_name = self.idx_to_class[top_class_id]
            top_confidence = float(top_prob)

            all_probs = {
                self.idx_to_class[i]: float(probs[i])
                for i in range(self.num_classes)
            }

            return top_class_name, top_class_id, top_confidence, all_probs


class EndpointHandler:
    def __init__(self, path: str = ""):
        """
        Initialize the handler by loading all models.
        Args:
            path: Path to the model directory (provided by HF)
        """
        self.device = "cuda" if torch.cuda.is_available() else "cpu"

        # YOLO Classification models
        self.models = {
            "business": YOLO(f"{path}/models/business_best.pt"),
            "finishing": YOLO(f"{path}/models/finishing_best.pt")
        }

        # Road Damage Detection model (YOLO)
        self.rdd_model = YOLO(f"{path}/models/rdd/yolo12s_RDD2022_best.pt")

        # SurfaceAI models (EfficientNetV2)
        self.surfaceai_models = {
            "surface_type": EfficientNetClassifier(f"{path}/models/surfaceai/surface_type_v1.pt", self.device),
            "road_type": EfficientNetClassifier(f"{path}/models/surfaceai/road_type_v1.pt", self.device),
            "quality": {
                "asphalt": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_asphalt_v1.pt", self.device),
                "concrete": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_concrete_v1.pt", self.device),
                "paving_stones": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_paving_stones_v1.pt", self.device),
                "sett": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_sett_v1.pt", self.device),
                "unpaved": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_unpaved_v1.pt", self.device),
            }
        }

    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_efficientnet(self, model: EfficientNetClassifier, image: Image.Image) -> Dict[str, Any]:
        """Run EfficientNet classification and return formatted results."""
        class_name, class_id, confidence, all_probs = model.predict(image)
        return {
            "class": class_name,
            "class_id": class_id,
            "confidence": round(confidence, 4),
            "all_probs": {k: round(v, 4) for k, v in all_probs.items()}
        }

    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_efficientnet(
            self.surfaceai_models["surface_type"], image
        )
        results["surface_type"] = surface_result

        # Get road type
        road_result = self._run_efficientnet(
            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_efficientnet(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_efficientnet(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)}]

	"""
	Hugging Face Inference Endpoint Custom Handler
	Handles inference for multiple models:
	- business/finishing: YOLO classification models
	- rdd: YOLO road damage detection (object detection with bounding boxes)
	- surfaceai: EfficientNetV2 models for surface type, road type, and quality classification
	"""

	import base64
	import io
	from typing import Any, Dict, List, Tuple
	from PIL import Image
	from ultralytics import YOLO
	import torch
	import torch.nn as nn
	import torch.nn.functional as F
	from torchvision import transforms
	from torchvision.models import efficientnet_v2_s


	class EfficientNetClassifier:
	"""Wrapper for EfficientNetV2 classification models."""

	def __init__(self, model_path: str, device: str = "cpu"):
	checkpoint = torch.load(model_path, map_location=device, weights_only=False)
	self.class_to_idx = checkpoint["class_to_idx"]
	self.idx_to_class = {v: k for k, v in self.class_to_idx.items()}
	self.num_classes = len(self.class_to_idx)
	self.is_regression = checkpoint.get("is_regression", False)
	self.device = device

	# Determine output size
	output_size = 1 if self.is_regression else self.num_classes

	# Build model
	self.model = efficientnet_v2_s(weights=None)
	self.model.classifier = nn.Sequential(
	nn.Dropout(p=0.2, inplace=True),
	nn.Linear(self.model.classifier[1].in_features, output_size)
	)
	self.model.load_state_dict(checkpoint["model_state_dict"])
	self.model.to(device)
	self.model.eval()

	# Image transforms (EfficientNetV2-S uses 384x384)
	self.transform = transforms.Compose([
	transforms.Resize((384, 384)),
	transforms.ToTensor(),
	transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
	])

	def predict(self, image: Image.Image) -> Tuple[str, int, float, Dict[str, float]]:
	"""Run inference and return class, id, confidence, and all probabilities."""
	image = image.convert("RGB")
	input_tensor = self.transform(image).unsqueeze(0).to(self.device)

	with torch.no_grad():
	outputs = self.model(input_tensor)

	if self.is_regression:
	# Regression model: output is a quality score
	raw_score = float(outputs[0, 0])
	# Clamp to valid range based on this model's class indices
	min_idx = min(self.idx_to_class.keys())
	max_idx = max(self.idx_to_class.keys())
	score = max(min_idx, min(max_idx, raw_score))
	class_id = int(round(score))
	# Ensure class_id is valid
	if class_id not in self.idx_to_class:
	class_id = min(self.idx_to_class.keys(), key=lambda x: abs(x - score))
	class_name = self.idx_to_class[class_id]

	# Create pseudo-probabilities based on distance from score
	all_probs = {}
	for idx, name in self.idx_to_class.items():
	distance = abs(idx - score)
	all_probs[name] = max(0, 1 - distance * 0.25)

	return class_name, class_id, score, all_probs
	else:
	# Classification model
	probs = F.softmax(outputs, dim=1)[0]
	top_prob, top_idx = torch.max(probs, 0)
	top_class_id = int(top_idx)
	top_class_name = self.idx_to_class[top_class_id]
	top_confidence = float(top_prob)

	all_probs = {
	self.idx_to_class[i]: float(probs[i])
	for i in range(self.num_classes)
	}

	return top_class_name, top_class_id, top_confidence, all_probs


	class EndpointHandler:
	def __init__(self, path: str = ""):
	"""
	Initialize the handler by loading all models.

	Args:
	path: Path to the model directory (provided by HF)
	"""
	self.device = "cuda" if torch.cuda.is_available() else "cpu"

	# YOLO Classification models
	self.models = {
	"business": YOLO(f"{path}/models/business_best.pt"),
	"finishing": YOLO(f"{path}/models/finishing_best.pt")
	}

	# Road Damage Detection model (YOLO)
	self.rdd_model = YOLO(f"{path}/models/rdd/yolo12s_RDD2022_best.pt")

	# SurfaceAI models (EfficientNetV2)
	self.surfaceai_models = {
	"surface_type": EfficientNetClassifier(f"{path}/models/surfaceai/surface_type_v1.pt", self.device),
	"road_type": EfficientNetClassifier(f"{path}/models/surfaceai/road_type_v1.pt", self.device),
	"quality": {
	"asphalt": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_asphalt_v1.pt", self.device),
	"concrete": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_concrete_v1.pt", self.device),
	"paving_stones": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_paving_stones_v1.pt", self.device),
	"sett": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_sett_v1.pt", self.device),
	"unpaved": EfficientNetClassifier(f"{path}/models/surfaceai/quality/surface_quality_unpaved_v1.pt", self.device),
	}
	}

	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_efficientnet(self, model: EfficientNetClassifier, image: Image.Image) -> Dict[str, Any]:
	"""Run EfficientNet classification and return formatted results."""
	class_name, class_id, confidence, all_probs = model.predict(image)
	return {
	"class": class_name,
	"class_id": class_id,
	"confidence": round(confidence, 4),
	"all_probs": {k: round(v, 4) for k, v in all_probs.items()}
	}

	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_efficientnet(
	self.surfaceai_models["surface_type"], image
	)
	results["surface_type"] = surface_result

	# Get road type
	road_result = self._run_efficientnet(
	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_efficientnet(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_efficientnet(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)}]