Spaces:
Sleeping
Sleeping
File size: 4,496 Bytes
2fb1271 73b3130 cdd2881 73b3130 cdd2881 a7a7037 cdd2881 73b3130 9270c11 73b3130 9270c11 2fb1271 cdd2881 73b3130 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 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 |
# 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()
|