Update segmentation.py

segmentation.py

@@ -69,19 +69,26 @@ def get_blurred_mask(img: Image, body_part_id: int):
     return dilated_mask_blurred
 
 
-def get_cropped_face(pil_image : Image):
-    image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_BGR2RGB)
+def get_cropped_face(pil_image: Image):
+    face = get_cropped(pil_image, 11, False)
+    image = np.array(face)
     face_casc = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-    
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+    gray = cv2.cvtColor(image, cv2.COLOR_RGBA2GRAY)
     faces = face_casc.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
+
+    if len(faces) == 0:
+        return pil_image
     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")
+    cropped_face = image[y:y+h, x:x+w]
+
+    result = Image.new('RGBA', pil_image.size, (255, 255, 255, 0))
+    face_pil = Image.fromarray(cropped_face)
+    if face_pil.size != (w, h):
+        face_pil = face_pil.resize((w, h))
+
+    result.paste(face_pil, (x, y))
+
+    return result