Spaces:
Running
Running
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)) | |
# Save the dilated mask for debugging if needed. | |
Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg") | |
# --- Extract outlines (only used for DXF generation) --- | |
t = time.time() | |
outlines, contours = extract_outlines(dilated_mask) | |
print("Outline extraction completed in {:.2f} seconds".format(time.time() - t)) | |
# Instead of drawing the original contours, we now prepare a clean copy of the shrunk image for drawing new contours. | |
output_img = shrunked_img.copy() | |
del shrunked_img | |
gc.collect() | |
# --- Generate DXF using the extracted contours and apply finger clearance --- | |
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 (Centered horizontally) --- | |
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) | |
# --- Draw only the new contours (final_polygons_inch) on the clean output image --- | |
draw_polygons_inch(final_polygons_inch, output_img, scaling_factor, processed_size[0], color=(0,0,255), thickness=2) | |
# Also prepare an "Outlines" image based on a blank canvas for clarity. | |
new_outlines = np.ones_like(output_img) * 255 | |
draw_polygons_inch(final_polygons_inch, new_outlines, 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(output_img, annotation_text.strip(), (text_px, text_py), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2, cv2.LINE_AA) | |
cv2.putText(new_outlines, annotation_text.strip(), (text_px, text_py), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2, cv2.LINE_AA) | |
outlines_color = cv2.cvtColor(new_outlines, cv2.COLOR_BGR2RGB) | |
print("Total prediction time: {:.2f} seconds".format(time.time() - overall_start)) | |
return ( | |
cv2.cvtColor(output_img, 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=[ | |
["./Test20.jpg", 0.075, "No", "No", 300.0, 200.0, "mm", "MyTool"], | |
["./Test21.jpg", 0.075, "Yes", "Yes", 300.0, 200.0, "mm", "Tool2"] | |
] | |
) | |
iface.launch(share=True) |