Spaces:
Running
on
Zero
Running
on
Zero
File size: 4,500 Bytes
d87a301 2db3b8e d87a301 2db3b8e d87a301 2db3b8e d87a301 8f50378 d87a301 |
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 |
import gradio as gr
import torch
from PIL import Image, ImageFilter, ImageDraw
from ultralytics import YOLO
import spaces
MODEL_PATH = "./models/face_yolov8n_v2.pt"
MODEL = YOLO(MODEL_PATH)
def create_rounded_rectangle_mask(image:Image.Image, radius, alpha=255):
size = image.size
factor = 5 # Factor to increase the image size that I can later antialiaze the corners
radius = radius * factor
image = Image.new('RGBA', (size[0] * factor, size[1] * factor), (0, 0, 0, 0))
# create corner
corner = Image.new('RGBA', (radius, radius), (0, 0, 0, 0))
draw = ImageDraw.Draw(corner)
# added the fill = .. you only drew a line, no fill
draw.pieslice((0, 0, radius * 2, radius * 2), 180, 270, fill=(255, 255, 255, alpha))
# max_x, max_y
mx, my = (size[0] * factor, size[1] * factor)
# paste corner rotated as needed
# use corners alpha channel as mask
image.paste(corner, (0, 0), corner)
image.paste(corner.rotate(90), (0, my - radius), corner.rotate(90))
image.paste(corner.rotate(180), (mx - radius, my - radius), corner.rotate(180))
image.paste(corner.rotate(270), (mx - radius, 0), corner.rotate(270))
# draw both inner rects
draw = ImageDraw.Draw(image)
draw.rectangle([(radius, 0), (mx - radius, my)], fill=(255, 255, 255, alpha))
draw.rectangle([(0, radius), (mx, my - radius)], fill=(255, 255, 255, alpha))
return image.resize(size, Image.LANCZOS) # Smooth the corners
def max_rounding_radius(rect_coords):
# Extract the coordinates of the rectangle [x0, y0, x1, y1]
rect_coords = [coord
-1 for coord in rect_coords]
x0, y0, x1, y1 = rect_coords
# Calculate the width and height of the rectangle
width = x1 - x0
height = y1 - y0
# Determine the smallest dimension (width or height)
min_dimension = min(width, height)
# Calculate the maximum radius as half of the smallest dimension
return min_dimension // 2
@spaces.GPU(enable_queue=True)
def generate_image(source_image:Image.Image, confidence=0.3, radius=50, blur_amount=10, margin=0):
if source_image is None:
return source_image
device = "cuda" if torch.cuda.is_available() else "cpu"
pred = MODEL(source_image, conf=confidence, device=device)
bboxes = pred[0].boxes.xyxy.cpu().numpy()
if bboxes.size == 0:
result = None
else:
bboxes = bboxes.tolist()
result = source_image.copy()
for bboxe in bboxes:
bboxe = [round(coord) for coord in bboxe]
mean_margin = margin // 2
new_bboxe = bboxe[0]-mean_margin, bboxe[1]-mean_margin, bboxe[2]+mean_margin, bboxe[3]+mean_margin
new_x0, new_y0, new_x1, new_y1 = new_bboxe
new_width = new_x1 - new_x0
new_height = new_y1 - new_y0
if not (new_width > source_image.width or new_height > source_image.height):
bboxe = new_bboxe
# Crop the region of interest
region = result.crop(bboxe)
final_radius = round(max_rounding_radius(bboxe)* (radius/100))
mask = create_rounded_rectangle_mask(region, final_radius)
# Apply blur filter to the cropped region
blurred_region = region.filter(ImageFilter.GaussianBlur(radius=blur_amount))
# Paste the blurred region back onto the original image
result.paste(blurred_region, bboxe, mask)
return result
css = """
img {
max-height: 500px;
object-fit: contain;
}
"""
with gr.Blocks(css=css) as FACE_2_BLUR:
with gr.Row():
with gr.Column():
image = gr.Image(label="Upload your image", type="pil")
generate_btn = gr.Button("Generate")
confidence = gr.Slider(label="Detection model confidence threshold.", value=0.3, minimum=0.0, maximum=1.0, step=0.01)
radius = gr.Slider(label="Edges radius in percentage.", value=50, minimum=0, maximum=100, step=1)
blur_amount = gr.Number(label="Controls the strength of the blur effect.", value=10, minimum=1, maximum=20, step=1)
margin = gr.Slider(label="Margin to add to the blurred box.", value=0, minimum=0, maximum=200)
with gr.Column():
image_out = gr.Image(label="Blurred face image", type="pil")
generate_btn.click(
fn=generate_image,
inputs=[image, confidence, radius, blur_amount, margin],
outputs=image_out,
api_name="generate_blurred"
)
FACE_2_BLUR.launch() |