Create app.py

app.py

@@ -0,0 +1,120 @@
+
+import gradio as gr
+import torch
+from PIL import Image, ImageFilter, ImageDraw
+from ultralytics import YOLO
+
+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
+
+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()