Spaces:
Sleeping
Sleeping
File size: 1,703 Bytes
36cd99b 267073f 36cd99b 267073f 36cd99b |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 |
"""
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']
|