"""
Created By: Ishwor Subedi
Date: 2024-07-17
"""
import cv2
import torchvision.transforms as T
import torch
import torchvision.models as models
from PIL import Image
from transformers import pipeline


class FaceShapeDetection:
    def __init__(self):
        self.model = models.efficientnet_b4(pretrained=False, num_classes=5)
        # Classes:[{'Heart': 0, 'Oblong': 1, 'Oval': 2, 'Round': 3, 'Square': 4}]
        self.classes = ['Heart', 'Oblong', 'Oval', 'Round', 'Square']
        self.model.load_state_dict(torch.load("/home/ishwor/Desktop/MannequinToModel/artifacts/best_model.pth"))
        self.model.eval()
        self.model.to("cpu")

        self.transform = T.Compose([
            T.ToPILImage(),
            T.Resize((224, 224)),
            T.ToTensor(),
            T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
        ])

    def detect_face_shape(self, image):
        image = self.transform(image)
        image = image.unsqueeze(0).to("cuda")

        with torch.no_grad():
            output = self.model(image)
        de_normalized_output = torch.nn.functional.softmax(output, dim=1)
        return self.classes[
            torch.argmax(de_normalized_output).item()], de_normalized_output[0][
            torch.argmax(de_normalized_output)].tolist()


class FaceShapeDetectionTransformer:
    def __init__(self):
        self.t_model = pipeline("image-classification", model="metadome/face_shape_classification",
                                torch_dtype=torch.float16,
                                device="cpu")

    def detect_face_shape(self, image):
        result = self.t_model(image)
        return result[0]['label'], result[0]['score']