Upload Models.py with huggingface_hub

dd9d3f9 verified almost 2 years ago

17.1 kB

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",
}

	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",
	}