Update app.py

app.py

@@ -38,19 +38,70 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
+# import gradio as gr
+# import cv2
+# import numpy as np
+# from ultralytics import YOLO
+
+
+# model = YOLO(r"best.pt")
+
+# def process_image(image):
+#     image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  
+#     results = model.predict(image, conf=0.15)
+
+#     detected_faces = len(results[0].boxes.cls) if results[0].boxes is not None else 0
+
+#     if detected_faces == 1:
+#         mask_tensor = results[0].masks.data[0].cpu().numpy()
+#         mask = (mask_tensor * 255).astype(np.uint8)
+        
+#         mask = cv2.resize(mask, (image.shape[1], image.shape[0]))
+#         kernel = np.ones((5, 5), np.uint8)
+#         mask = cv2.dilate(mask, kernel, iterations=2)
+#         mask = cv2.erode(mask, kernel, iterations=2)
+        
+#         rgba_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGBA)
+#         rgba_image[:, :, 3] = mask
+        
+#         return rgba_image
+    
+#     # Return a blank image with error text
+#     error_image = np.zeros((300, 500, 3), dtype=np.uint8)
+#     error_message = "Error: More than one face detected."
+#     cv2.putText(error_image, error_message, (20, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2, cv2.LINE_AA)
+    
+#     return error_image
+
+# demo = gr.Interface(
+#     fn=process_image,
+#     inputs=gr.Image(type="numpy"),
+#     outputs=gr.Image(type="numpy"),
+#     title="Face Segmentation",
+#     description="Upload a singale face image"
+# )
+
+# if __name__ == "__main__":
+#     demo.launch()
+
+
+
 import gradio as gr
 import cv2
 import numpy as np
 from ultralytics import YOLO
 
 
-model = YOLO(r"best.pt")
+face_model = YOLO(r"face_det.pt")
+seg_model = YOLO(r"seg_model.pt")
 
 def process_image(image):
     image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)  
-    results = model.predict(image, conf=0.15)
+    face_results = face_model.predict(image, conf=0.5, save=True)
+    results = seg_model.predict(image, conf=0.15)
 
-    detected_faces = len(results[0].boxes.cls) if results[0].boxes is not None else 0
+    detected_faces = len(face_results[0].boxes.cls) if results[0].boxes is not None else 0
+    print(detected_faces)
 
     if detected_faces == 1:
         mask_tensor = results[0].masks.data[0].cpu().numpy()
@@ -65,13 +116,21 @@ def process_image(image):
         rgba_image[:, :, 3] = mask
         
         return rgba_image
+    elif detected_faces >= 1:
+        # Return a blank image with error text
+        error_image = np.zeros((300, 500, 3), dtype=np.uint8)
+        error_message = "Error: More than one face detected."
+        cv2.putText(error_image, error_message, (20, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2, cv2.LINE_AA)
+        return error_image
+    elif detected_faces == 0:
+        # Return a blank image with error text
+        error_image = np.zeros((300, 500, 3), dtype=np.uint8)
+        error_message = "Error: No face detected."
+        cv2.putText(error_image, error_message, (20, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2, cv2.LINE_AA)
+        return error_image
+
     
-    # Return a blank image with error text
-    error_image = np.zeros((300, 500, 3), dtype=np.uint8)
-    error_message = "Error: More than one face detected."
-    cv2.putText(error_image, error_message, (20, 150), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2, cv2.LINE_AA)
     
-    return error_image
 
 demo = gr.Interface(
     fn=process_image,