Spaces:
Sleeping
Sleeping
File size: 2,554 Bytes
2fb1271 cdd2881 9270c11 2fb1271 9270c11 2fb1271 cdd2881 9270c11 b5f073e 9270c11 cdd2881 0f97f8b cdd2881 9270c11 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
# 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.Please 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,
inputs=gr.Image(type="numpy"),
outputs=gr.Image(type="numpy"),
title="Face Segmentation",
description="Upload a singale face image"
)
if __name__ == "__main__":
demo.launch()
|