# *************************************************************************** #
#                                                                              #
#    prediction.py                                                             #
#                                                                              #
#    By: Widium <ebennace@student.42lausanne.ch>                               #
#    Github : https://github.com/widium                                        #
#                                                                              #
#    Created: 2023/04/25 14:48:03 by Widium                                    #
#    Updated: 2023/04/25 14:48:03 by Widium                                    #
#                                                                              #
# **************************************************************************** #

from torch import Tensor
from PIL import Image
from torchvision.transforms import Compose


def img_to_tensor(pil_image : Image, img_transformer : Compose)->Tensor:
    """
    Preprocess the input image to a batched tensor 

    Args:
        pil_image (Image): input image open with PIL
        img_transformer (Compose): same image transformer used for train the model 

    Returns:
        Tensor: tensor batched of input img 
    """
    img_tensor = img_transformer(pil_image)
    img_tensor_batched = img_tensor.unsqueeze(dim=0)
    return (img_tensor_batched)

# =============================================================================== #

from torch import Tensor
from typing import List, Dict

def associate_probs_with_class_names(
    probs : Tensor,
    class_names : List[str]
)->Dict[str, float]:
    """
    Create dictionary with 'class_name' : probs with Softmax output Tensor

    Args:
        probs (Tensor): Softmax output tensor
        class_names (List[str]): list of class names

    Returns:
        Dict[str, float]: dictionary who asscociate class name as key with probs as value
    """
    
    predictions = dict()
    probs_array = probs.squeeze()
    
    for index, prob in enumerate(probs_array):
        
        class_name = class_names[index]
        prob = round(float(prob), 2)
        
        predictions[f"{class_name}"] = prob
    
    return (predictions)