Upload Models.py with huggingface_hub

Models.py

@@ -0,0 +1,457 @@
+import os
+# import sys
+from functools import partial
+from pathlib import Path
+import logging
+
+import torch
+from huggingface_hub import hf_hub_download
+from torch import Tensor, nn
+from torchvision import models, transforms
+import pandas as pd
+from collections import defaultdict
+
+
+class ModelInterface:
+    """
+    Interface for managing image classification and regression tasks.
+
+    """
+    def __init__(self, config):
+        """
+        Initialize the ModelInterface.
+
+        Parameters:
+            config (dict): Configuration dictionary containing the following keys:
+                - gpu_kernel (int): GPU index to use for computations. Defaults to the first available GPU if available, otherwise CPU.
+                - transform_surface (dict): Parameters for surface type and quality image transformations, including resize, crop, and normalization settings.
+                - transform_road_type (dict): Parameters for road type image transformations, similar to surface transformations.
+                - model_root (str): Directory path where model files are stored locally. Defaults to folder name 'models'.
+                - models (dict): Dictionary mapping prediction levels (e.g., 'road_type', 'surface_type') to model file names.
+                - hf_model_repo (str): Hugging Face repository ID for downloading models if not found locally.
+        """
+        self.device = self._validate_device(config.get('gpu_kernel', ''))
+        self.model_root = Path(config.get("model_root", "models"))
+        self.models = config.get("models")
+        self.hf_model_repo = config.get("hf_model_repo", "")
+        self._validate_models()
+        self._default_normalization = (NORM_MEAN, NORM_SD)
+        self.transform_surface = self._validate_transform(config.get("transform_surface", None), "surface_type")
+        self.transform_road_type = self._validate_transform(config.get("transform_road_type", None), "road_type")
+
+    def _validate_device(self, gpu_kernel):
+        try:
+            cuda = "cuda" if gpu_kernel == '' else f"cuda:{gpu_kernel}"
+            return torch.device(
+                cuda if torch.cuda.is_available() else "cpu"
+            )
+        except Exception as e:
+            logging.warning(f"An unexpected error occurred while selecting GPU: {e}\n"
+                          + "Falling back to CPU.")
+            return torch.device("cpu")
+
+    def _validate_models(self):
+        """
+        Check if model files exist and download from hugging face if not.
+        """
+        if self.models is None:
+            raise TypeError("No models are defined.")
+
+        log_model_not_defined = "No model for '{level_string}' is defined. Prediction is skipped."
+
+        # check surface type model
+        level = "surface_type"
+        model_file = self.models.get(level)
+        if model_file is None:
+            logging.warning(log_model_not_defined.format(level_string=model_to_info_string[level]))
+        else:
+            self.download_model(model_file)
+            _, surface_class_to_idx, _ = self.load_model(model=model_file)
+
+        # check quality models
+        level = "surface_quality"
+        sub_models = self.models.get(level)
+        if model_file is None:
+            logging.warning(log_model_not_defined.format(level_string=model_to_info_string[level]))
+        else:
+            for surface_type in surface_class_to_idx:
+                model_file = sub_models.get(surface_type)
+                if model_file is None:
+                    logging.warning(log_model_not_defined.format(level_string=surface_type))
+                else:
+                    self.download_model(model_file)
+                    self.load_model(model=model_file)
+
+        # check road type model
+        level = "road_type"
+        model_file = self.models.get(level)
+        if model_file is None:
+            logging.warning(log_model_not_defined.format(level_string=model_to_info_string[level]))
+        else:
+            self.download_model(model_file)
+            self.load_model(model=model_file)
+
+    def _validate_transform(self, transform, level):
+        """
+        Validate the transformation for a given model type if the model exists.
+
+        Parameters:
+            - transform (dict): transformation.
+            - level (str): model level.
+
+        Returns:
+            dict: transformation.
+        """
+        if (level in self.models) and (transform is None):
+            logging.warning(f"No transformation for {model_to_info_string[level]} prediction defined.")
+            transform = {}
+        
+        if "normalize" not in transform:
+            logging.info(f"No normalization parameters for {model_to_info_string[level]} prediction provided. Using default values.")
+            transform["normalize"] = self._default_normalization
+        
+        return transform
+  
+
+    def download_model(self, model):
+        """
+        Download a model from Hugging Face repository.
+
+        Parameters:
+            - model (str): Model file name.
+
+        Returns:
+            None
+        """
+        model_path = self.model_root / model
+        # load model data from hugging face if not locally available
+        if not os.path.exists(model_path):
+            logging.info(
+                f"Model file not found at {model_path}. Downloading from Hugging Face..."
+            )
+            try:
+                os.makedirs(self.model_root, exist_ok=True)
+                model_path = hf_hub_download(
+                    repo_id=self.hf_model_repo, filename=model, local_dir=self.model_root
+                )
+                logging.info(f"Model file downloaded successfully to {model_path}.")
+            except Exception as e:
+                logging.error(f"An unexpected error occurred while downloading the model {model}: {e}")
+                raise e
+
+
+    @staticmethod
+    def custom_crop(img, crop_style=None):
+        """
+        Crop an image according to the specified style.
+
+        Parameters:
+            - img (PIL.Image): Input image to be cropped.
+            - crop_style (str, optional): Style of cropping (e.g., 'lower_middle_half').
+
+        Returns:
+            PIL.Image: Cropped image.
+        """
+        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:  # None, or not valid
+            logging.warning(f"Cropping method {crop_style} is not defined. Image is not cropped.")
+            return img
+
+        cropped_img = transforms.functional.crop(img, top, left, height, width)
+        return cropped_img
+
+    def transform(
+        self,
+        resize=None,
+        crop=None,
+        to_tensor=True,
+        normalize=None,
+    ):
+        """
+        Create a PyTorch image transformation function based on specified parameters.
+
+        Parameters:
+            - resize ((int, int) or int, optional): Target size for resizing, e.g. (height, width). If int, then used for both height and width.
+            - crop (str, optional): crop style e.g. 'lower_middle_third'
+            - to_tensor (bool, optional): Converts the PIL Image (H x W x C) in the range [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0]
+            - normalize (tuple of lists [r, g, b], optional): Mean and standard deviation for normalization.
+
+        Returns:
+            PyTorch image transformation function.
+        """
+        transform_list = []
+
+        if crop is not None:
+            transform_list.append(
+                transforms.Lambda(partial(self.custom_crop, crop_style=crop))
+            )
+
+        if resize is not None:
+            if isinstance(resize, int):
+                resize = (resize, resize)
+            transform_list.append(transforms.Resize(resize))
+
+        if to_tensor:
+            transform_list.append(transforms.ToTensor())
+
+        if normalize is not None:
+            transform_list.append(transforms.Normalize(*normalize))
+
+        composed_transform = transforms.Compose(transform_list)
+        return composed_transform
+
+    def preprocessing(self, img_data_raw, transform):
+        """
+        Preprocess raw image data using a specified transformation.
+
+        Parameters:
+            - img_data_raw (list): List of raw images to preprocess.
+            - transform (dict): Dictionary of transformation parameters.
+
+        Returns:
+            torch.Tensor: Preprocessed image tensor.
+        """
+        if not img_data_raw:
+            raise ValueError("Image data is empty.")
+        
+        transform = self.transform(**transform)
+        img_data = torch.stack([transform(img) for img in img_data_raw])
+        return img_data
+
+    def load_model(self, model):
+        """
+        Load a model from local storage.
+
+        Parameters:
+            - model (str): Model file name.
+
+        Returns:
+            nn.Module: Loaded model.
+            dict: Mapping of classes to indices.
+            bool: Whether the model is for regression.
+        """
+        model_path = self.model_root / model
+        try:
+            model_state = torch.load(model_path, map_location=self.device)
+            model_name = model_state["model_name"]
+            is_regression = model_state["is_regression"]
+            class_to_idx = model_state["class_to_idx"]
+            num_classes = 1 if is_regression else len(class_to_idx.items())
+            model_state_dict = model_state["model_state_dict"]
+            model_cls = model_mapping[model_name]
+            model = model_cls(num_classes=num_classes)
+            model.load_state_dict(model_state_dict)
+        except Exception as e:
+            logging.error(f"An unexpected error occurred while loading the model {model_path}: {e}")
+            raise e
+
+        return model, class_to_idx, is_regression
+
+    def predict(self, model, data):
+        """
+        Perform predictions using the specified model and input data.
+
+        Parameters:
+            - model (nn.Module): The model to use for predictions.
+            - data (torch.Tensor): Batch of input data.
+
+        Returns:
+            torch.Tensor: Predicted values or class probabilities.
+        """
+        model.to(self.device)
+        model.eval()
+
+        image_batch = data.to(self.device)
+
+        with torch.no_grad():
+            batch_outputs = model(image_batch)
+            # batch_classes, batch_values = model.get_class_and_value(batch_outputs)
+            batch_values = model.get_class_probabilities(batch_outputs)
+
+        return batch_values
+    
+    @staticmethod
+    def predict_to_classes(batch_values, class_to_idx):
+        """
+        Map predicted values to classes.
+
+        Parameters:
+            - batch_values (torch.Tensor): Batch of prediction values.
+            - class_to_idx (dict): Mapping from class names to indices.
+
+        Returns:
+            list: List of predicted values.
+            list: List of predicted classes.
+        """
+        idx_to_class = {i: cls for cls, i in class_to_idx.items()}
+        
+        if len(list(batch_values.shape)) < 2:
+            classes = [
+                idx_to_class[
+                    min(
+                        max(idx.item(), min(list(class_to_idx.values()))),
+                        max(list(class_to_idx.values())),
+                    )
+                ]
+                for idx in batch_values.round().int()
+            ]
+            values = batch_values.tolist()
+        else:
+            classes = [idx_to_class[idx.item()] for idx in torch.argmax(batch_values, dim=1)]
+            values = batch_values.tolist()
+            
+        return values, classes
+
+    def batch_classifications(self, img_data_raw, img_ids=None):
+        """
+        Perform batch classification for multiple prediction levels (road type, surface type, surface quality).
+
+        Parameters:
+            - img_data_raw (list): List of raw images to classify.
+            - img_ids (list, optional): List of IDs corresponding to the images. Defaults to indices.
+
+        Returns:
+            list: Combined list of image ids and predictions across levels.
+        """
+        if not img_data_raw:
+            logging.info("Input data is empty. No predictions performed.")
+            return []
+        
+        # default image ids
+        if img_ids is None:
+            img_ids = range(len(img_data_raw))
+
+        # road type
+        level = "road_type"
+        model_file = self.models.get(level)
+        if model_file is not None:
+            model, class_to_idx, _ = self.load_model(model=model_file)
+            data = self.preprocessing(img_data_raw, self.transform_road_type)
+            values = self.predict(model, data)
+            road_values, road_classes = self.predict_to_classes(values, class_to_idx)
+
+        # surface type
+        level = "surface_type"
+        model_file = self.models.get(level)
+        if model_file is not None:
+            model, class_to_idx, _ = self.load_model(model=model_file)
+            data = self.preprocessing(img_data_raw, self.transform_surface)
+            values = self.predict(model, data)
+            surface_values, surface_classes = self.predict_to_classes(values, class_to_idx)
+
+            # surface quality
+            level = "surface_quality"
+            sub_models = self.models.get(level)
+            if sub_models is not None:
+                surface_indices = defaultdict(list)
+                for i, surface_type in enumerate(surface_classes):
+                    surface_indices[surface_type].append(i)
+
+                quality_values = [None] * len(img_data_raw)
+                quality_classes = [None] * len(img_data_raw)
+                for surface_type, indices in surface_indices.items():
+                    model_file = sub_models.get(surface_type)
+                    if model_file is not None:
+                        model, class_to_idx, _ = self.load_model(model=model_file)
+                        values = self.predict(model, data[indices])
+                        values, classes = self.predict_to_classes(values, class_to_idx)
+                        for idx, vl, cls in zip(indices, values, classes):
+                            quality_values[idx] = vl
+                            quality_classes[idx] = cls
+
+        # final results combination
+        final_results = [
+            [
+                img_ids[i],
+                road_classes[i],
+                road_values[i],
+                surface_classes[i],
+                surface_values[i],
+                quality_classes[i],
+                quality_values[i],
+            ]
+            for i in range(len(img_data_raw))
+        ]
+
+        return final_results
+
+
+class CustomEfficientNetV2SLinear(nn.Module):
+    """
+    Custom implementation of EfficientNetV2-S with a linear classifier for classification or regression tasks.
+
+    Attributes:
+        features (nn.Sequential): Feature extractor from EfficientNetV2-S.
+        avgpool (nn.AdaptiveAvgPool2d): Adaptive average pooling layer.
+        classifier (nn.Sequential): Fully connected layers for classification.
+        is_regression (bool): Whether the model is configured for regression tasks.
+        criterion (callable): Loss function used for training the model.
+    """
+
+    def __init__(self, num_classes, avg_pool=1):
+        super(CustomEfficientNetV2SLinear, self).__init__()
+
+        model = models.efficientnet_v2_s(weights="IMAGENET1K_V1")
+        # adapt output layer
+        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
+        if num_classes == 1:
+            self.criterion = nn.MSELoss
+            self.is_regression = True
+        else:
+            self.criterion = nn.CrossEntropyLoss
+            self.is_regression = False
+
+    def get_class_probabilities(self, x):
+        if self.is_regression:
+            x = x.flatten()
+        else:
+            x = nn.functional.softmax(x, dim=1)
+        return x
+
+    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_optimizer_layers(self):
+    #     return self.classifier
+
+
+# Constants
+EFFNET_LINEAR = "efficientNetV2SLinear"
+CROP_LOWER_MIDDLE_HALF = "lower_middle_half"
+CROP_LOWER_HALF = "lower_half"
+NORM_MEAN = [0.42834484577178955, 0.4461250305175781, 0.4350937306880951]
+NORM_SD = [0.22991590201854706, 0.23555299639701843, 0.26348039507865906]
+
+
+model_mapping = {
+    EFFNET_LINEAR: CustomEfficientNetV2SLinear,
+}
+
+model_to_info_string = {
+    "surface_type": "surface type",
+    "road_type": "road type",
+    "surface_quality": "quality",
+}