# 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)
    
#     if len(results[0].boxes.cls) == 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
#     else:
#         return "Error: Uploaded image has more than one face. Please upload a different image."


# demo = gr.Interface(
#     fn=process_image,
#     inputs=gr.Image(type="numpy"),
#     outputs=gr.Image(type="numpy"),
#     title="Face Segmentation",
#     description="Upload an 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")

# 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


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)  
    face_results = face_model.predict(image, conf=0.5, save=True)
    results = seg_model.predict(image, conf=0.15)

    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()
        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
    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

    
    

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()