Create app.py

app.py

@@ -0,0 +1,578 @@
+from __future__ import annotations
+import os
+import gc
+import base64
+import io
+import time
+import shutil
+import numpy as np
+import torch
+import cv2
+import ezdxf
+import gradio as gr
+from PIL import Image, ImageEnhance
+from pathlib import Path
+from typing import List, Union
+from ultralytics import YOLOWorld, YOLO
+from ultralytics.engine.results import Results
+from ultralytics.utils.plotting import save_one_box
+from transformers import AutoModelForImageSegmentation
+from torchvision import transforms
+from scalingtestupdated import calculate_scaling_factor
+from shapely.geometry import Polygon, Point, MultiPolygon
+from scipy.interpolate import splprep, splev
+from scipy.ndimage import gaussian_filter1d
+from u2net import U2NETP
+
+# ---------------------
+# Create a cache folder for models
+# ---------------------
+CACHE_DIR = os.path.join(os.path.dirname(__file__), ".cache")
+os.makedirs(CACHE_DIR, exist_ok=True)
+
+# ---------------------
+# Custom Exceptions
+# ---------------------
+class DrawerNotDetectedError(Exception):
+    """Raised when the drawer cannot be detected in the image"""
+    pass
+
+class ReferenceBoxNotDetectedError(Exception):
+    """Raised when the reference box cannot be detected in the image"""
+    pass
+
+# ---------------------
+# Global Model Initialization with caching and print statements
+# ---------------------
+print("Loading YOLOWorld model...")
+start_time = time.time()
+yolo_model_path = os.path.join(CACHE_DIR, "yolov8x-worldv2.pt")
+if not os.path.exists(yolo_model_path):
+    print("Caching YOLOWorld model to", yolo_model_path)
+    shutil.copy("yolov8x-worldv2.pt", yolo_model_path)
+drawer_detector_global = YOLOWorld(yolo_model_path)
+drawer_detector_global.set_classes(["box"])
+print("YOLOWorld model loaded in {:.2f} seconds".format(time.time() - start_time))
+
+print("Loading YOLO reference model...")
+start_time = time.time()
+reference_model_path = os.path.join(CACHE_DIR, "last.pt")
+if not os.path.exists(reference_model_path):
+    print("Caching YOLO reference model to", reference_model_path)
+    shutil.copy("last.pt", reference_model_path)
+reference_detector_global = YOLO(reference_model_path)
+print("YOLO reference model loaded in {:.2f} seconds".format(time.time() - start_time))
+
+print("Loading U²-Net model for reference background removal (U2NETP)...")
+start_time = time.time()
+u2net_model_path = os.path.join(CACHE_DIR, "u2netp.pth")
+if not os.path.exists(u2net_model_path):
+    print("Caching U²-Net model to", u2net_model_path)
+    shutil.copy("u2netp.pth", u2net_model_path)
+u2net_global = U2NETP(3, 1)
+u2net_global.load_state_dict(torch.load(u2net_model_path, map_location="cpu"))
+device = "cpu"
+u2net_global.to(device)
+u2net_global.eval()
+print("U²-Net model loaded in {:.2f} seconds".format(time.time() - start_time))
+
+print("Loading BiRefNet model...")
+start_time = time.time()
+birefnet_global = AutoModelForImageSegmentation.from_pretrained(
+    "zhengpeng7/BiRefNet", trust_remote_code=True, cache_dir=CACHE_DIR
+)
+torch.set_float32_matmul_precision("high")
+birefnet_global.to(device)
+birefnet_global.eval()
+print("BiRefNet model loaded in {:.2f} seconds".format(time.time() - start_time))
+
+# Define transform for BiRefNet
+transform_image_global = transforms.Compose([
+    transforms.Resize((1024, 1024)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+])
+
+# ---------------------
+# Model Reload Function (if needed)
+# ---------------------
+def unload_and_reload_models():
+    global drawer_detector_global, reference_detector_global, birefnet_global, u2net_global
+    print("Reloading models...")
+    start_time = time.time()
+    del drawer_detector_global, reference_detector_global, birefnet_global, u2net_global
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+    gc.collect()
+    new_drawer_detector = YOLOWorld(os.path.join(CACHE_DIR, "yolov8x-worldv2.pt"))
+    new_drawer_detector.set_classes(["box"])
+    new_reference_detector = YOLO(os.path.join(CACHE_DIR, "last.pt"))
+    new_birefnet = AutoModelForImageSegmentation.from_pretrained(
+        "zhengpeng7/BiRefNet", trust_remote_code=True, cache_dir=CACHE_DIR
+    )
+    new_birefnet.to(device)
+    new_birefnet.eval()
+    new_u2net = U2NETP(3, 1)
+    new_u2net.load_state_dict(torch.load(os.path.join(CACHE_DIR, "u2netp.pth"), map_location="cpu"))
+    new_u2net.to(device)
+    new_u2net.eval()
+    drawer_detector_global = new_drawer_detector
+    reference_detector_global = new_reference_detector
+    birefnet_global = new_birefnet
+    u2net_global = new_u2net
+    print("Models reloaded in {:.2f} seconds".format(time.time() - start_time))
+
+# ---------------------
+# Helper Function: resize_img (defined once)
+# ---------------------
+def resize_img(img: np.ndarray, resize_dim):
+    return np.array(Image.fromarray(img).resize(resize_dim))
+
+# ---------------------
+# Other Helper Functions for Detection & Processing
+# ---------------------
+def yolo_detect(image: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor]) -> np.ndarray:
+    t = time.time()
+    results: List[Results] = drawer_detector_global.predict(image)
+    if not results or len(results) == 0 or len(results[0].boxes) == 0:
+        raise DrawerNotDetectedError("Drawer not detected in the image.")
+    print("Drawer detection completed in {:.2f} seconds".format(time.time() - t))
+    return save_one_box(results[0].cpu().boxes.xyxy, im=results[0].orig_img, save=False)
+
+def detect_reference_square(img: np.ndarray):
+    t = time.time()
+    res = reference_detector_global.predict(img, conf=0.45)
+    if not res or len(res) == 0 or len(res[0].boxes) == 0:
+        raise ReferenceBoxNotDetectedError("Reference box not detected in the image.")
+    print("Reference detection completed in {:.2f} seconds".format(time.time() - t))
+    return (
+        save_one_box(res[0].cpu().boxes.xyxy, res[0].orig_img, save=False),
+        res[0].cpu().boxes.xyxy[0]
+    )
+
+# Use U2NETP for reference background removal.
+def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:
+    t = time.time()
+    image_pil = Image.fromarray(image)
+    transform_u2netp = transforms.Compose([
+        transforms.Resize((320, 320)),
+        transforms.ToTensor(),
+        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+    ])
+    input_tensor = transform_u2netp(image_pil).unsqueeze(0).to("cpu")
+    with torch.no_grad():
+        outputs = u2net_global(input_tensor)
+    pred = outputs[0]
+    pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
+    pred_np = pred.squeeze().cpu().numpy()
+    pred_np = cv2.resize(pred_np, (image_pil.width, image_pil.height))
+    pred_np = (pred_np * 255).astype(np.uint8)
+    print("U2NETP background removal completed in {:.2f} seconds".format(time.time() - t))
+    return pred_np
+
+# Use BiRefNet for main object background removal.
+def remove_bg(image: np.ndarray) -> np.ndarray:
+    t = time.time()
+    image_pil = Image.fromarray(image)
+    input_images = transform_image_global(image_pil).unsqueeze(0).to("cpu")
+    with torch.no_grad():
+        preds = birefnet_global(input_images)[-1].sigmoid().cpu()
+    pred = preds[0].squeeze()
+    pred_pil = transforms.ToPILImage()(pred)
+    scale_ratio = 1024 / max(image_pil.size)
+    scaled_size = (int(image_pil.size[0] * scale_ratio), int(image_pil.size[1] * scale_ratio))
+    result = np.array(pred_pil.resize(scaled_size))
+    print("BiRefNet background removal completed in {:.2f} seconds".format(time.time() - t))
+    return result
+
+def make_square(img: np.ndarray):
+    height, width = img.shape[:2]
+    max_dim = max(height, width)
+    pad_height = (max_dim - height) // 2
+    pad_width = (max_dim - width) // 2
+    pad_height_extra = max_dim - height - 2 * pad_height
+    pad_width_extra = max_dim - width - 2 * pad_width
+    if len(img.shape) == 3:
+        padded = np.pad(img, ((pad_height, pad_height + pad_height_extra),
+                              (pad_width, pad_width + pad_width_extra),
+                              (0, 0)), mode="edge")
+    else:
+        padded = np.pad(img, ((pad_height, pad_height + pad_height_extra),
+                              (pad_width, pad_width + pad_width_extra)), mode="edge")
+    return padded
+
+def shrink_bbox(image: np.ndarray, shrink_factor: float):
+    height, width = image.shape[:2]
+    center_x, center_y = width // 2, height // 2
+    new_width = int(width * shrink_factor)
+    new_height = int(height * shrink_factor)
+    x1 = max(center_x - new_width // 2, 0)
+    y1 = max(center_y - new_height // 2, 0)
+    x2 = min(center_x + new_width // 2, width)
+    y2 = min(center_y + new_height // 2, height)
+    return image[y1:y2, x1:x2]
+
+def exclude_scaling_box(image: np.ndarray, bbox: np.ndarray, orig_size: tuple, processed_size: tuple, expansion_factor: float = 1.2) -> np.ndarray:
+    x_min, y_min, x_max, y_max = map(int, bbox)
+    scale_x = processed_size[1] / orig_size[1]
+    scale_y = processed_size[0] / orig_size[0]
+    x_min = int(x_min * scale_x)
+    x_max = int(x_max * scale_x)
+    y_min = int(y_min * scale_y)
+    y_max = int(y_max * scale_y)
+    box_width = x_max - x_min
+    box_height = y_max - y_min
+    expanded_x_min = max(0, int(x_min - (expansion_factor - 1) * box_width / 2))
+    expanded_x_max = min(image.shape[1], int(x_max + (expansion_factor - 1) * box_width / 2))
+    expanded_y_min = max(0, int(y_min - (expansion_factor - 1) * box_height / 2))
+    expanded_y_max = min(image.shape[0], int(y_max + (expansion_factor - 1) * box_height / 2))
+    image[expanded_y_min:expanded_y_max, expanded_x_min:expanded_x_max] = 0
+    return image
+
+def resample_contour(contour):
+    num_points = 1000
+    smoothing_factor = 5
+    spline_degree = 3
+    if len(contour) < spline_degree + 1:
+        raise ValueError(f"Contour must have at least {spline_degree + 1} points, but has {len(contour)} points.")
+    contour = contour[:, 0, :]
+    tck, _ = splprep([contour[:, 0], contour[:, 1]], s=smoothing_factor)
+    u = np.linspace(0, 1, num_points)
+    resampled_points = splev(u, tck)
+    smoothed_x = gaussian_filter1d(resampled_points[0], sigma=1)
+    smoothed_y = gaussian_filter1d(resampled_points[1], sigma=1)
+    return np.array([smoothed_x, smoothed_y]).T
+
+# ---------------------
+# Add the missing extract_outlines function
+# ---------------------
+def extract_outlines(binary_image: np.ndarray) -> (np.ndarray, list):
+    contours, _ = cv2.findContours(binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+    outline_image = np.zeros_like(binary_image)
+    cv2.drawContours(outline_image, contours, -1, (255), thickness=2)
+    return cv2.bitwise_not(outline_image), contours
+
+# ---------------------
+# Functions for Finger Cut Clearance
+# ---------------------
+def union_tool_and_circle(tool_polygon: Polygon, center_inch, circle_diameter=1.0):
+    radius = circle_diameter / 2.0
+    circle_poly = Point(center_inch).buffer(radius, resolution=64)
+    union_poly = tool_polygon.union(circle_poly)
+    return union_poly
+
+def build_tool_polygon(points_inch):
+    return Polygon(points_inch)
+
+def polygon_to_exterior_coords(poly: Polygon):
+    if poly.geom_type == "MultiPolygon":
+        biggest = max(poly.geoms, key=lambda g: g.area)
+        poly = biggest
+    if not poly.exterior:
+        return []
+    return list(poly.exterior.coords)
+
+def place_finger_cut_randomly(tool_polygon, points_inch, existing_centers, all_polygons, circle_diameter=1.0, min_gap=0.25, max_attempts=30):
+    import random
+    needed_center_distance = circle_diameter + min_gap
+    radius = circle_diameter / 2.0
+    for _ in range(max_attempts):
+        idx = random.randint(0, len(points_inch) - 1)
+        cx, cy = points_inch[idx]
+        too_close = False
+        for (ex_x, ex_y) in existing_centers:
+            if np.hypot(cx - ex_x, cy - ex_y) < needed_center_distance:
+                too_close = True
+                break
+        if too_close:
+            continue
+        circle_poly = Point((cx, cy)).buffer(radius, resolution=64)
+        union_poly = tool_polygon.union(circle_poly)
+        overlap_with_others = False
+        too_close_to_others = False
+        for poly in all_polygons:
+            if union_poly.intersects(poly):
+                overlap_with_others = True
+                break
+            if circle_poly.buffer(min_gap).intersects(poly):
+                too_close_to_others = True
+                break
+        if overlap_with_others or too_close_to_others:
+            continue
+        existing_centers.append((cx, cy))
+        return union_poly, (cx, cy)
+    print("Warning: Could not place a finger cut circle meeting all spacing requirements.")
+    return None, None
+
+# ---------------------
+# DXF Spline and Boundary Functions
+# ---------------------
+def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=False):
+    degree = 3
+    closed = True
+    doc = ezdxf.new(units=0)
+    doc.units = ezdxf.units.IN
+    doc.header["$INSUNITS"] = ezdxf.units.IN
+    msp = doc.modelspace()
+    finger_cut_centers = []
+    final_polygons_inch = []
+    for contour in inflated_contours:
+        try:
+            resampled_contour = resample_contour(contour)
+            points_inch = [(x * scaling_factor, (height - y) * scaling_factor) for x, y in resampled_contour]
+            if len(points_inch) < 3:
+                continue
+            if np.linalg.norm(np.array(points_inch[0]) - np.array(points_inch[-1])) > 1e-2:
+                points_inch.append(points_inch[0])
+            tool_polygon = build_tool_polygon(points_inch)
+            if finger_clearance:
+                union_poly, center = place_finger_cut_randomly(tool_polygon, points_inch, finger_cut_centers, final_polygons_inch, circle_diameter=1.0, min_gap=0.25, max_attempts=30)
+                if union_poly is not None:
+                    tool_polygon = union_poly
+            exterior_coords = polygon_to_exterior_coords(tool_polygon)
+            if len(exterior_coords) < 3:
+                continue
+            msp.add_spline(exterior_coords, degree=degree, dxfattribs={"layer": "TOOLS"})
+            final_polygons_inch.append(tool_polygon)
+        except ValueError as e:
+            print(f"Skipping contour: {e}")
+    return doc, final_polygons_inch
+
+def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width, boundary_unit):
+    msp = doc.modelspace()
+    if boundary_unit == "mm":
+        boundary_length_in = boundary_length / 25.4
+        boundary_width_in = boundary_width / 25.4
+    else:
+        boundary_length_in = boundary_length
+        boundary_width_in = boundary_width
+    min_x = float("inf")
+    min_y = float("inf")
+    max_x = -float("inf")
+    max_y = -float("inf")
+    for poly in polygons_inch:
+        b = poly.bounds
+        min_x = min(min_x, b[0])
+        min_y = min(min_y, b[1])
+        max_x = max(max_x, b[2])
+        max_y = max(max_y, b[3])
+    if min_x == float("inf"):
+        print("No tool polygons found, skipping boundary.")
+        return None
+    shape_cx = (min_x + max_x) / 2
+    shape_cy = (min_y + max_y) / 2
+    half_w = boundary_width_in / 2.0
+    half_l = boundary_length_in / 2.0
+    left = shape_cx - half_w
+    right = shape_cx + half_w
+    bottom = shape_cy - half_l
+    top = shape_cy + half_l
+    rect_coords = [(left, bottom), (right, bottom), (right, top), (left, top), (left, bottom)]
+    from shapely.geometry import Polygon as ShapelyPolygon
+    boundary_polygon = ShapelyPolygon(rect_coords)
+    msp.add_lwpolyline(rect_coords, close=True, dxfattribs={"layer": "BOUNDARY"})
+    return boundary_polygon
+
+def draw_polygons_inch(polygons_inch, image_rgb, scaling_factor, image_height, color=(0,0,255), thickness=2):
+    for poly in polygons_inch:
+        if poly.geom_type == "MultiPolygon":
+            for subpoly in poly.geoms:
+                draw_single_polygon(subpoly, image_rgb, scaling_factor, image_height, color, thickness)
+        else:
+            draw_single_polygon(poly, image_rgb, scaling_factor, image_height, color, thickness)
+
+def draw_single_polygon(poly, image_rgb, scaling_factor, image_height, color=(0,0,255), thickness=2):
+    ext = list(poly.exterior.coords)
+    if len(ext) < 3:
+        return
+    pts_px = []
+    for (x_in, y_in) in ext:
+        px = int(x_in / scaling_factor)
+        py = int(image_height - (y_in / scaling_factor))
+        pts_px.append([px, py])
+    pts_px = np.array(pts_px, dtype=np.int32)
+    cv2.polylines(image_rgb, [pts_px], isClosed=True, color=color, thickness=thickness, lineType=cv2.LINE_AA)
+
+# ---------------------
+# Main Predict Function with Finger Cut Clearance, Boundary Box, Annotation and Sharpness Enhancement
+# ---------------------
+def predict(
+    image: Union[str, bytes, np.ndarray],
+    offset_inches: float,
+    finger_clearance: str,    # "Yes" or "No"
+    add_boundary: str,        # "Yes" or "No"
+    boundary_length: float,
+    boundary_width: float,
+    boundary_unit: str,
+    annotation_text: str
+):
+    overall_start = time.time()
+    # Convert image to NumPy array if needed.
+    if isinstance(image, str):
+        if os.path.exists(image):
+            image = np.array(Image.open(image).convert("RGB"))
+        else:
+            try:
+                image = np.array(Image.open(io.BytesIO(base64.b64decode(image))).convert("RGB"))
+            except Exception:
+                raise ValueError("Invalid base64 image data")
+    # Apply sharpness enhancement if image is a NumPy array.
+    if isinstance(image, np.ndarray):
+        pil_image = Image.fromarray(image)
+        enhanced_image = ImageEnhance.Sharpness(pil_image).enhance(1.5)
+        image = np.array(enhanced_image)
+    try:
+        t = time.time()
+        drawer_img = yolo_detect(image)
+        print("Drawer detection completed in {:.2f} seconds".format(time.time() - t))
+        t = time.time()
+        shrunked_img = make_square(shrink_bbox(drawer_img, 0.90))
+        del drawer_img
+        gc.collect()
+        print("Image shrinking completed in {:.2f} seconds".format(time.time() - t))
+    except DrawerNotDetectedError:
+        raise DrawerNotDetectedError("Drawer not detected! Please take another picture with a drawer.")
+    try:
+        t = time.time()
+        reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
+        print("Reference square detection completed in {:.2f} seconds".format(time.time() - t))
+    except ReferenceBoxNotDetectedError:
+        raise ReferenceBoxNotDetectedError("Reference box not detected! Please take another picture with a reference box.")
+    t = time.time()
+    reference_obj_img = make_square(reference_obj_img)
+    reference_square_mask = remove_bg_u2netp(reference_obj_img)
+    print("Reference image processing completed in {:.2f} seconds".format(time.time() - t))
+    t = time.time()
+    try:
+        cv2.imwrite("mask.jpg", cv2.cvtColor(reference_obj_img, cv2.COLOR_RGB2GRAY))
+        scaling_factor = calculate_scaling_factor(
+            reference_image_path="./Reference_ScalingBox.jpg",
+            target_image=reference_square_mask,
+            feature_detector="ORB",
+        )
+    except ZeroDivisionError:
+        scaling_factor = None
+        print("Error calculating scaling factor: Division by zero")
+    except Exception as e:
+        scaling_factor = None
+        print(f"Error calculating scaling factor: {e}")
+    if scaling_factor is None or scaling_factor == 0:
+        scaling_factor = 1.0
+        print("Using default scaling factor of 1.0 due to calculation error")
+    gc.collect()
+    print("Scaling factor determined: {}".format(scaling_factor))
+    t = time.time()
+    orig_size = shrunked_img.shape[:2]
+    objects_mask = remove_bg(shrunked_img)
+    processed_size = objects_mask.shape[:2]
+    objects_mask = exclude_scaling_box(objects_mask, scaling_box_coords, orig_size, processed_size, expansion_factor=1.2)
+    objects_mask = resize_img(objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0]))
+    del scaling_box_coords
+    gc.collect()
+    print("Object masking completed in {:.2f} seconds".format(time.time() - t))
+    t = time.time()
+    offset_pixels = (offset_inches / scaling_factor) * 2 + 1 if scaling_factor != 0 else 1
+    dilated_mask = cv2.dilate(objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8))
+    del objects_mask
+    gc.collect()
+    print("Mask dilation completed in {:.2f} seconds".format(time.time() - t))
+    Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
+    t = time.time()
+    outlines, contours = extract_outlines(dilated_mask)
+    shrunked_img_contours = cv2.drawContours(shrunked_img.copy(), contours, -1, (0, 0, 255), thickness=2)
+    del shrunked_img
+    gc.collect()
+    print("Outline extraction completed in {:.2f} seconds".format(time.time() - t))
+    t = time.time()
+    use_finger_clearance = True if finger_clearance.lower() == "yes" else False
+    doc, final_polygons_inch = save_dxf_spline(contours, scaling_factor, processed_size[0], finger_clearance=use_finger_clearance)
+    del contours
+    gc.collect()
+    print("DXF generation completed in {:.2f} seconds".format(time.time() - t))
+    boundary_polygon = None
+    if add_boundary.lower() == "yes":
+        boundary_polygon = add_rectangular_boundary(doc, final_polygons_inch, boundary_length, boundary_width, boundary_unit)
+        if boundary_polygon is not None:
+            final_polygons_inch.append(boundary_polygon)
+    # --- Annotation Text Placement (Bottom-Right) ---
+    min_x = float("inf")
+    min_y = float("inf")
+    max_x = -float("inf")
+    max_y = -float("inf")
+    for poly in final_polygons_inch:
+        b = poly.bounds
+        if b[0] < min_x:
+            min_x = b[0]
+        if b[1] < min_y:
+            min_y = b[1]
+        if b[2] > max_x:
+            max_x = b[2]
+        if b[3] > max_y:
+            max_y = b[3]
+    margin = 0.5
+    text_x = (min_x + max_x) / 2
+    text_y = min_y - margin
+    msp = doc.modelspace()
+    if annotation_text.strip():
+        text_entity = msp.add_text(
+            annotation_text.strip(),
+            dxfattribs={
+                "height": 0.25,
+                "layer": "ANNOTATION"
+            }
+        )
+        text_entity.dxf.insert = (text_x, text_y)
+    dxf_filepath = os.path.join("./outputs", "out.dxf")
+    doc.saveas(dxf_filepath)
+    # --- End Annotation Text Placement ---
+    draw_polygons_inch(final_polygons_inch, shrunked_img_contours, scaling_factor, processed_size[0], color=(0,0,255), thickness=2)
+    outlines_bgr = cv2.cvtColor(outlines, cv2.COLOR_GRAY2BGR)
+    draw_polygons_inch(final_polygons_inch, outlines_bgr, scaling_factor, processed_size[0], color=(0,0,255), thickness=2)
+    if annotation_text.strip():
+        text_px = int(text_x / scaling_factor)
+        text_py = int(processed_size[0] - (text_y / scaling_factor))
+        cv2.putText(shrunked_img_contours, annotation_text.strip(), (text_px, text_py), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2, cv2.LINE_AA)
+        cv2.putText(outlines_bgr, annotation_text.strip(), (text_px, text_py), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2, cv2.LINE_AA)
+    outlines_color = cv2.cvtColor(outlines_bgr, cv2.COLOR_BGR2RGB)
+    print("Total prediction time: {:.2f} seconds".format(time.time() - overall_start))
+    return (
+        cv2.cvtColor(shrunked_img_contours, cv2.COLOR_BGR2RGB),
+        outlines_color,
+        dxf_filepath,
+        dilated_mask,
+        str(scaling_factor)
+    )
+
+# ---------------------
+# Gradio Interface
+# ---------------------
+if __name__ == "__main__":
+    os.makedirs("./outputs", exist_ok=True)
+    def gradio_predict(img, offset, finger_clearance, add_boundary, boundary_length, boundary_width, boundary_unit, annotation_text):
+        return predict(img, offset, finger_clearance, add_boundary, boundary_length, boundary_width, boundary_unit, annotation_text)
+    iface = gr.Interface(
+        fn=gradio_predict,
+        inputs=[
+            gr.Image(label="Input Image"),
+            gr.Number(label="Offset value for Mask (inches)", value=0.075),
+            gr.Dropdown(label="Add Finger Clearance?", choices=["Yes", "No"], value="No"),
+            gr.Dropdown(label="Add Rectangular Boundary?", choices=["Yes", "No"], value="No"),
+            gr.Number(label="Boundary Length", value=300.0, precision=2),
+            gr.Number(label="Boundary Width", value=200.0, precision=2),
+            gr.Dropdown(label="Boundary Unit", choices=["mm", "inches"], value="mm"),
+            gr.Textbox(label="Annotation (max 20 chars)", max_length=20, placeholder="Type up to 20 characters")
+        ],
+        outputs=[
+            gr.Image(label="Output Image"),
+            gr.Image(label="Outlines of Objects"),
+            gr.File(label="DXF file"),
+            gr.Image(label="Mask"),
+            gr.Textbox(label="Scaling Factor (inches/pixel)")
+        ],
+        examples=[
+            ["./examples/Test20.jpg", 0.075, "No", "No", 300.0, 200.0, "mm", "MyTool"],
+            ["./examples/Test21.jpg", 0.075, "Yes", "Yes", 300.0, 200.0, "mm", "Tool2"]
+        ]
+    )
+    iface.launch(share=True)