Update app.py

app.py

@@ -38,7 +38,6 @@
 # if __name__ == "__main__":
 #     demo.launch()
 
-
 import gradio as gr
 import cv2
 import numpy as np
@@ -50,9 +49,9 @@ 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)
@@ -66,10 +65,13 @@ def process_image(image):
         rgba_image[:, :, 3] = mask
         
         return rgba_image
-    elif detected_faces > 1:
-        return "Error: Uploaded image has more than one face. Please upload a different image."
-    else:
-        return "Error: No face detected. Please upload a valid 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.\nPlease upload a different image."
+    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,
@@ -81,3 +83,4 @@ demo = gr.Interface(
 
 if __name__ == "__main__":
     demo.launch()
+