Update segmentation.py

segmentation.py

@@ -4,6 +4,8 @@ from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation,
 import numpy as np
 import torch.nn as nn
 from scipy.ndimage import binary_dilation
+import cv2
+
 
 model = None
 extractor = None
@@ -14,13 +16,9 @@ def init_body():
     model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes").to("cuda")
 
 
-def get_mask(img: Image, body_part_id: int, inverse=False, face=False):
-    if face:
-        inputs = extractor_face(images=img, return_tensors="pt").to("cuda")
-        outputs = model_face(**inputs)
-    else:
-        inputs = extractor(images=img, return_tensors="pt").to("cuda")
-        outputs = model(**inputs)
+def get_mask(img: Image, body_part_id: int, inverse=False):
+    inputs = extractor(images=img, return_tensors="pt").to("cuda")
+    outputs = model(**inputs)
     logits = outputs.logits.cpu()
 
     upsampled_logits = nn.functional.interpolate(
@@ -40,9 +38,9 @@ def get_mask(img: Image, body_part_id: int, inverse=False, face=False):
     return pil_seg
 
 
-def get_cropped(img: Image, body_part_id: int, inverse:bool, face:bool):
+def get_cropped(img: Image, body_part_id: int, inverse:bool):
     
-    pil_seg = get_mask(img, body_part_id, inverse, face)
+    pil_seg = get_mask(img, body_part_id, inverse)
     crop_mask_np = np.array(pil_seg.convert('L'))
     crop_mask_binary = crop_mask_np > 128
 
@@ -71,4 +69,19 @@ def get_blurred_mask(img: Image, body_part_id: int):
     return dilated_mask_blurred
 
 
+def get_face_crop(pil_image : Image):
+    image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_BGR2RGB)
+    face_casc = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
+    
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    faces = face_casc.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
+    x, y, w, h = faces[0]
+    
+    cropped_face =  np.ones_like(np.array(pil_image))*255
+    cropped_face[y:y+h, x:x+w] = image[y:y+h, x:x+w]
+    
+    img = cv2.cvtColor(cropped_face, cv2.COLOR_BGR2RGB)
+    return Image.fromarray(img.astype("uint8")).convert("RGB")
+
+