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