Spaces:
Running
Running
| # 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 | |
| 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." | |
| 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() | |