import os from pathlib import Path from typing import List, Union, Tuple from PIL import Image import ezdxf.units import numpy as np import torch from torchvision import transforms from ultralytics import YOLOWorld, YOLO from ultralytics.engine.results import Results from ultralytics.utils.plotting import save_one_box from transformers import AutoModelForImageSegmentation import cv2 import ezdxf import gradio as gr import gc # from scalingtestupdated import calculate_scaling_factor from scalingtestupdated import calculate_scaling_factor_with_units, calculate_paper_scaling_factor, convert_units, calculate_paper_scaling_factor_corrected from scipy.interpolate import splprep, splev from scipy.ndimage import gaussian_filter1d import json import time import signal from shapely.ops import unary_union from shapely.geometry import MultiPolygon, GeometryCollection, Polygon, Point from u2netp import U2NETP import logging import shutil import sys # Add this at the very beginning of your main Python file, before any other imports os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '0' os.environ['OPENCV_IO_ENABLE_JASPER'] = '0' os.environ['QT_QPA_PLATFORM'] = 'offscreen' os.environ['MPLBACKEND'] = 'Agg' # For headless environments import matplotlib matplotlib.use('Agg') # Initialize logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Create cache directory for models CACHE_DIR = os.path.join(os.path.dirname(__file__), ".cache") os.makedirs(CACHE_DIR, exist_ok=True) # Paper size configurations (in mm) PAPER_SIZES = { "A4": {"width": 210, "height": 297}, "A3": {"width": 297, "height": 420}, "US Letter": {"width": 215.9, "height": 279.4} } # Custom Exception Classes class TimeoutReachedError(Exception): pass class BoundaryOverlapError(Exception): pass class TextOverlapError(Exception): pass class PaperNotDetectedError(Exception): """Raised when the paper cannot be detected in the image""" pass class MultipleObjectsError(Exception): """Raised when multiple objects are detected on the paper""" def __init__(self, message="Multiple objects detected. Please place only a single object on the paper."): super().__init__(message) class NoObjectDetectedError(Exception): """Raised when no object is detected on the paper""" def __init__(self, message="No object detected on the paper. Please ensure an object is placed on the paper."): super().__init__(message) class FingerCutOverlapError(Exception): """Raised when finger cuts overlap with existing geometry""" def __init__(self, message="There was an overlap with fingercuts... Please try again to generate dxf."): super().__init__(message) class ReferenceBoxNotDetectedError(Exception): """Raised when reference box/paper cannot be detected""" def __init__(self, message="Reference box not detected"): super().__init__(message) # Global model variables for lazy loading paper_detector_global = None u2net_global = None birefnet = None # Model paths paper_model_path = os.path.join(CACHE_DIR, "paper_detector.pt") # You'll need to train/provide this u2net_model_path = os.path.join(CACHE_DIR, "u2netp.pth") # Global variable for YOLOWorld yolo_v8_global = None yolo_v8_model_path = os.path.join(CACHE_DIR, "yolov8n.pt") # Adjust path as needed # Device configuration device = "cpu" torch.set_float32_matmul_precision(["high", "highest"][0]) def ensure_model_files(): """Ensure model files are available in cache directory""" if not os.path.exists(paper_model_path): if os.path.exists("paper_detector.pt"): shutil.copy("paper_detector.pt", paper_model_path) else: logger.warning("paper_detector.pt model file not found - using fallback detection") if not os.path.exists(u2net_model_path): if os.path.exists("u2netp.pth"): shutil.copy("u2netp.pth", u2net_model_path) else: raise FileNotFoundError("u2netp.pth model file not found") logger.info("YOLOv8 will auto-download if not present") ensure_model_files() # Lazy loading functions def get_paper_detector(): """Lazy load paper detector model""" global paper_detector_global if paper_detector_global is None: logger.info("Loading paper detector model...") if os.path.exists(paper_model_path): try: paper_detector_global = YOLO(paper_model_path) logger.info("Paper detector loaded successfully") except Exception as e: logger.error(f"Failed to load paper detector: {e}") paper_detector_global = None else: # Fallback to generic object detection for paper-like rectangles logger.warning("Paper model file not found, using fallback detection") paper_detector_global = None return paper_detector_global def get_yolo_v8(): """Lazy load YOLOv8 model""" global yolo_v8_global if yolo_v8_global is None: logger.info("Loading YOLOv8 model...") try: yolo_v8_global = YOLO(yolo_v8_model_path) # Auto-downloads if needed logger.info("YOLOv8 model loaded successfully") except Exception as e: logger.error(f"Failed to load YOLOv8: {e}") yolo_v8_global = None return yolo_v8_global def get_u2net(): """Lazy load U2NETP model""" global u2net_global if u2net_global is None: logger.info("Loading U2NETP model...") u2net_global = U2NETP(3, 1) u2net_global.load_state_dict(torch.load(u2net_model_path, map_location="cpu")) u2net_global.to(device) u2net_global.eval() logger.info("U2NETP model loaded successfully") return u2net_global def load_birefnet_model(): """Load BiRefNet model from HuggingFace""" return AutoModelForImageSegmentation.from_pretrained( 'ZhengPeng7/BiRefNet', trust_remote_code=True ) def get_birefnet(): """Lazy load BiRefNet model""" global birefnet if birefnet is None: logger.info("Loading BiRefNet model...") birefnet = load_birefnet_model() birefnet.to(device) birefnet.eval() logger.info("BiRefNet model loaded successfully") return birefnet def detect_paper_contour(image: np.ndarray, output_unit: str = "mm") -> Tuple[np.ndarray, float]: """ Detect paper in the image using contour detection as fallback Returns the paper contour and estimated scaling factor """ logger.info("Using contour-based paper detection") # Convert to grayscale gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image # Apply bilateral filter to reduce noise while preserving edges filtered = cv2.bilateralFilter(gray, 9, 75, 75) # Apply adaptive threshold thresh = cv2.adaptiveThreshold(filtered, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 11, 2) # Edge detection with multiple thresholds edges1 = cv2.Canny(filtered, 50, 150) edges2 = cv2.Canny(filtered, 30, 100) edges = cv2.bitwise_or(edges1, edges2) # Morphological operations to connect broken edges kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel) # Find contours contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # Filter contours by area and aspect ratio to find paper-like rectangles paper_contours = [] image_area = image.shape[0] * image.shape[1] min_area = image_area * 0.20 # At least 15% of image max_area = image_area * 0.85 # At most 95% of image for contour in contours: area = cv2.contourArea(contour) if min_area < area < max_area: # Approximate contour to polygon epsilon = 0.015 * cv2.arcLength(contour, True) approx = cv2.approxPolyDP(contour, epsilon, True) # Check if it's roughly rectangular (4 corners) or close to it if len(approx) >= 4: # Calculate bounding rectangle rect = cv2.boundingRect(approx) w, h = rect[2], rect[3] aspect_ratio = w / h if h > 0 else 0 # Check if aspect ratio matches common paper ratios # A4: 1.414, A3: 1.414, US Letter: 1.294 if 1.3 < aspect_ratio < 1.5: # More lenient tolerance # Check if contour area is close to bounding rect area (rectangularity) rect_area = w * h if rect_area > 0: extent = area / rect_area if extent > 0.85: # At least 85% rectangular paper_contours.append((contour, area, aspect_ratio, extent)) if not paper_contours: logger.error("No paper-like contours found") raise ReferenceBoxNotDetectedError("Could not detect paper in the image using contour detection") # Select the best paper contour based on area and rectangularity paper_contours.sort(key=lambda x: (x[1] * x[3]), reverse=True) # Sort by area * extent best_contour = paper_contours[0][0] logger.info(f"Paper detected using contours: area={paper_contours[0][1]}, aspect_ratio={paper_contours[0][2]:.2f}") # Return 0.0 as placeholder - will be calculated later based on paper size return best_contour, 0.0 def detect_paper_bounds(image: np.ndarray, paper_size: str, output_unit: str = "mm") -> Tuple[np.ndarray, float]: """ Detect paper bounds in the image and calculate scaling factor """ try: paper_detector = get_paper_detector() if paper_detector is not None: # Use trained model if available results = paper_detector.predict(image, conf=0.8, verbose=False) if not results or len(results) == 0: logger.warning("No results from paper detector") return detect_paper_contour(image, output_unit) # Check if boxes exist and are not empty if not hasattr(results[0], 'boxes') or results[0].boxes is None or len(results[0].boxes) == 0: logger.warning("No boxes detected by model, using fallback contour detection") return detect_paper_contour(image, output_unit) # Get the largest detected paper boxes = results[0].boxes.xyxy.cpu().numpy() if len(boxes) == 0: logger.warning("Empty boxes detected, using fallback") return detect_paper_contour(image, output_unit) largest_box = None max_area = 0 for box in boxes: x_min, y_min, x_max, y_max = box area = (x_max - x_min) * (y_max - y_min) if area > max_area: max_area = area largest_box = box if largest_box is None: logger.warning("No valid paper box found, using fallback") return detect_paper_contour(image, output_unit) # Convert box to contour-like format x_min, y_min, x_max, y_max = map(int, largest_box) paper_contour = np.array([ [[x_min, y_min]], [[x_max, y_min]], [[x_max, y_max]], [[x_min, y_max]] ]) logger.info(f"Paper detected by model: {x_min},{y_min} to {x_max},{y_max}") else: # Use fallback contour detection logger.info("Using fallback contour detection for paper") paper_contour, _ = detect_paper_contour(image, output_unit) # After getting paper_contour, expand it rect = cv2.boundingRect(paper_contour) expansion = int(min(rect[2], rect[3]) * 0.1) # Expand by 10% x, y, w, h = rect expanded_contour = np.array([ [[max(0, x - expansion), max(0, y - expansion)]], [[min(image.shape[1], x + w + expansion), max(0, y - expansion)]], [[min(image.shape[1], x + w + expansion), min(image.shape[0], y + h + expansion)]], [[max(0, x - expansion), min(image.shape[0], y + h + expansion)]] ]) paper_contour = expanded_contour # Calculate scaling factor based on paper size with proper units # scaling_factor = calculate_paper_scaling_factor(paper_contour, paper_size, output_unit) scaling_factor, unit_string = calculate_paper_scaling_factor_corrected( paper_contour, paper_size, output_unit="mm", correction_factor=1.235, # Adjust this value method="average" # Try different methods ) return paper_contour, scaling_factor except Exception as e: logger.error(f"Error in paper detection: {e}") raise ReferenceBoxNotDetectedError(f"Failed to detect paper: {str(e)}") def calculate_paper_scaling_factor(paper_contour: np.ndarray, paper_size: str, output_unit: str = "mm") -> float: """ Calculate scaling factor based on detected paper dimensions with proper unit handling. """ from scalingtestupdated import calculate_paper_scaling_factor as calc_paper_scale scaling_factor, unit_string = calc_paper_scale(paper_contour, paper_size, output_unit) return scaling_factor def validate_single_object(mask: np.ndarray, paper_contour: np.ndarray) -> None: """ Validate that only a single object is present on the paper """ # Create a mask for the paper area paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8) cv2.fillPoly(paper_mask, [paper_contour], 255) # Apply paper mask to object mask masked_objects = cv2.bitwise_and(mask, paper_mask) # Find contours of objects within paper bounds contours, _ = cv2.findContours(masked_objects, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) # Filter out very small contours (noise) and paper-sized contours image_area = mask.shape[0] * mask.shape[1] min_area = 100 # Minimum area threshold max_area = image_area * 0.5 # Maximum 50% of image area (to exclude paper detection) significant_contours = [c for c in contours if min_area < cv2.contourArea(c) < max_area] if len(significant_contours) == 0: raise NoObjectDetectedError() elif len(significant_contours) > 1: raise MultipleObjectsError() logger.info(f"Single object validated: {len(significant_contours)} significant contour(s) found") def remove_bg_u2netp(image: np.ndarray) -> np.ndarray: """Remove background using U2NETP model""" try: u2net_model = get_u2net() 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(device) with torch.no_grad(): outputs = u2net_model(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) return pred_np except Exception as e: logger.error(f"Error in U2NETP background removal: {e}") raise def remove_bg(image: np.ndarray) -> np.ndarray: """Remove background using BiRefNet model for main objects""" try: birefnet_model = get_birefnet() transform_image = transforms.Compose([ transforms.Resize((1024, 1024)), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) image_pil = Image.fromarray(image) input_images = transform_image(image_pil).unsqueeze(0).to(device) with torch.no_grad(): preds = birefnet_model(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)) return np.array(pred_pil.resize(scaled_size)) except Exception as e: logger.error(f"Error in BiRefNet background removal: {e}") raise def mask_paper_area_in_image(image: np.ndarray, paper_contour: np.ndarray) -> np.ndarray: """Less aggressive masking to preserve corner objects""" masked_image = image.copy() # Much less aggressive shrinking - only 2% instead of 8% rect = cv2.boundingRect(paper_contour) shrink_pixels = max(5, int(min(rect[2], rect[3]) * 0.02)) # Changed from 0.08 to 0.02 x, y, w, h = rect # Create mask but keep more area outer_mask = np.ones(image.shape[:2], dtype=np.uint8) * 255 inner_contour = np.array([ [[x + shrink_pixels, y + shrink_pixels]], [[x + w - shrink_pixels, y + shrink_pixels]], [[x + w - shrink_pixels, y + h - shrink_pixels]], [[x + shrink_pixels, y + h - shrink_pixels]] ]) cv2.fillPoly(outer_mask, [inner_contour], 0) masked_image[outer_mask == 255] = [128, 128, 128] # Gray instead of black return masked_image def exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.2) -> np.ndarray: """Less aggressive paper area exclusion""" # Create paper mask paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8) cv2.fillPoly(paper_mask, [paper_contour], 255) # Instead of eroding, slightly expand the paper mask rect = cv2.boundingRect(paper_contour) expansion = max(10, int(min(rect[2], rect[3]) * 0.02)) # 2% expansion kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (expansion, expansion)) expanded_paper_mask = cv2.dilate(paper_mask, kernel, iterations=1) # Keep objects within expanded paper area result_mask = cv2.bitwise_and(mask, expanded_paper_mask) return result_mask def resample_contour(contour, edge_radius_px: int = 0): """Resample contour with radius-aware smoothing and periodic handling.""" logger.info(f"Starting resample_contour with contour of shape {contour.shape}") num_points = 1500 sigma = max(2, int(edge_radius_px) // 4) if len(contour) < 4: error_msg = f"Contour must have at least 4 points, but has {len(contour)} points." logger.error(error_msg) raise ValueError(error_msg) try: contour = contour[:, 0, :] logger.debug(f"Reshaped contour to shape {contour.shape}") if not np.array_equal(contour[0], contour[-1]): contour = np.vstack([contour, contour[0]]) tck, u = splprep(contour.T, u=None, s=0, per=True) u_new = np.linspace(u.min(), u.max(), num_points) x_new, y_new = splev(u_new, tck, der=0) if sigma > 0: x_new = gaussian_filter1d(x_new, sigma=sigma, mode='wrap') y_new = gaussian_filter1d(y_new, sigma=sigma, mode='wrap') x_new[-1] = x_new[0] y_new[-1] = y_new[0] result = np.array([x_new, y_new]).T logger.info(f"Completed resample_contour with result shape {result.shape}") return result except Exception as e: logger.error(f"Error in resample_contour: {e}") raise def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=False): """Save contours as DXF splines with optional finger cuts - scaling_factor should be in mm/px""" doc = ezdxf.new(units=ezdxf.units.MM) doc.header["$INSUNITS"] = ezdxf.units.MM msp = doc.modelspace() final_polygons_mm = [] # Use mm instead of inch for clarity finger_centers = [] original_polygons = [] for contour in inflated_contours: try: resampled_contour = resample_contour(contour) # Convert pixel coordinates to mm using the scaling factor points_mm = [(x * scaling_factor, (height - y) * scaling_factor) for x, y in resampled_contour] if len(points_mm) < 3: continue tool_polygon = build_tool_polygon(points_mm) original_polygons.append(tool_polygon) if finger_clearance: try: tool_polygon, center = place_finger_cut_adjusted( tool_polygon, points_mm, finger_centers, final_polygons_mm ) except FingerCutOverlapError: tool_polygon = original_polygons[-1] exterior_coords = polygon_to_exterior_coords(tool_polygon) if len(exterior_coords) < 3: continue # Coordinates are already in mm, so add directly to DXF msp.add_spline(exterior_coords, degree=3, dxfattribs={"layer": "TOOLS"}) final_polygons_mm.append(tool_polygon) except ValueError as e: logger.warning(f"Skipping contour: {e}") dxf_filepath = os.path.join("./outputs", "out.dxf") doc.saveas(dxf_filepath) return dxf_filepath, final_polygons_mm, original_polygons def build_tool_polygon(points_inch): """Build a polygon from inch-converted points""" return Polygon(points_inch) def polygon_to_exterior_coords(poly): """Extract exterior coordinates from polygon""" logger.info(f"Starting polygon_to_exterior_coords with input geometry type: {poly.geom_type}") try: if poly.geom_type == "GeometryCollection" or poly.geom_type == "MultiPolygon": logger.debug(f"Performing unary_union on {poly.geom_type}") unified = unary_union(poly) if unified.is_empty: logger.warning("unary_union produced an empty geometry; returning empty list") return [] if unified.geom_type == "GeometryCollection" or unified.geom_type == "MultiPolygon": largest = None max_area = 0.0 for g in getattr(unified, "geoms", []): if hasattr(g, "area") and g.area > max_area and hasattr(g, "exterior"): max_area = g.area largest = g if largest is None: logger.warning("No valid Polygon found in unified geometry; returning empty list") return [] poly = largest else: poly = unified if not hasattr(poly, "exterior") or poly.exterior is None: logger.warning("Input geometry has no exterior ring; returning empty list") return [] raw_coords = list(poly.exterior.coords) total = len(raw_coords) logger.info(f"Extracted {total} raw exterior coordinates") if total == 0: return [] # Subsample coordinates to at most 100 points max_pts = 100 if total > max_pts: step = total // max_pts sampled = [raw_coords[i] for i in range(0, total, step)] if sampled[-1] != raw_coords[-1]: sampled.append(raw_coords[-1]) logger.info(f"Downsampled perimeter from {total} to {len(sampled)} points") return sampled else: return raw_coords except Exception as e: logger.error(f"Error in polygon_to_exterior_coords: {e}") return [] def place_finger_cut_adjusted( tool_polygon: Polygon, points_inch: list, existing_centers: list, all_polygons: list, circle_diameter: float = 25.4, min_gap: float = 0.5, max_attempts: int = 100 ) -> Tuple[Polygon, tuple]: """Place finger cuts with collision avoidance""" logger.info(f"Starting place_finger_cut_adjusted with {len(points_inch)} input points") def fallback_solution(): logger.warning("Using fallback approach for finger cut placement") fallback_center = points_inch[len(points_inch) // 2] r = circle_diameter / 2.0 fallback_circle = Point(fallback_center).buffer(r, resolution=32) try: union_poly = tool_polygon.union(fallback_circle) except Exception as e: logger.warning(f"Fallback union failed ({e}); trying buffer-union fallback") union_poly = tool_polygon.buffer(0).union(fallback_circle.buffer(0)) existing_centers.append(fallback_center) logger.info(f"Fallback finger cut placed at {fallback_center}") return union_poly, fallback_center r = circle_diameter / 2.0 needed_center_dist = circle_diameter + min_gap raw_perimeter = polygon_to_exterior_coords(tool_polygon) if not raw_perimeter: logger.warning("No valid exterior coords found; using fallback immediately") return fallback_solution() if len(raw_perimeter) > 100: step = len(raw_perimeter) // 100 perimeter_coords = raw_perimeter[::step] logger.info(f"Subsampled perimeter from {len(raw_perimeter)} to {len(perimeter_coords)} points") else: perimeter_coords = raw_perimeter[:] indices = list(range(len(perimeter_coords))) np.random.shuffle(indices) logger.debug(f"Shuffled perimeter indices for candidate order") start_time = time.time() timeout_secs = 5.0 attempts = 0 try: while attempts < max_attempts: if time.time() - start_time > timeout_secs - 0.1: logger.warning(f"Approaching timeout after {attempts} attempts") return fallback_solution() for idx in indices: if time.time() - start_time > timeout_secs - 0.05: logger.warning("Timeout during candidate-point loop") return fallback_solution() cx, cy = perimeter_coords[idx] for dx, dy in [(0, 0), (-min_gap/2, 0), (min_gap/2, 0), (0, -min_gap/2), (0, min_gap/2)]: candidate_center = (cx + dx, cy + dy) # Check distance to existing finger centers too_close_finger = any( np.hypot(candidate_center[0] - ex, candidate_center[1] - ey) < needed_center_dist for (ex, ey) in existing_centers ) if too_close_finger: continue # Build candidate circle candidate_circle = Point(candidate_center).buffer(r, resolution=32) # Must overlap ≥30% with this polygon try: inter_area = tool_polygon.intersection(candidate_circle).area except Exception: continue if inter_area < 0.3 * candidate_circle.area: continue # Must not intersect other polygons invalid = False for other_poly in all_polygons: if other_poly.equals(tool_polygon): continue if other_poly.buffer(min_gap).intersects(candidate_circle) or \ other_poly.buffer(min_gap).touches(candidate_circle): invalid = True break if invalid: continue # Union and return try: union_poly = tool_polygon.union(candidate_circle) if union_poly.geom_type == "MultiPolygon" and len(union_poly.geoms) > 1: continue if union_poly.equals(tool_polygon): continue except Exception: continue existing_centers.append(candidate_center) logger.info(f"Finger cut placed successfully at {candidate_center} after {attempts} attempts") return union_poly, candidate_center attempts += 1 if attempts >= (max_attempts // 2) and (time.time() - start_time) > timeout_secs * 0.8: logger.warning(f"Approaching timeout (attempt {attempts})") return fallback_solution() logger.warning(f"No valid spot after {max_attempts} attempts, using fallback") return fallback_solution() except Exception as e: logger.error(f"Error in place_finger_cut_adjusted: {e}") return fallback_solution() # def extract_outlines(binary_image: np.ndarray) -> Tuple[np.ndarray, list]: # """Extract outlines from binary image""" # contours, _ = cv2.findContours( # binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE # ) # outline_image = np.full_like(binary_image, 255) # return outline_image, contours def extract_outlines(binary_image: np.ndarray) -> Tuple[np.ndarray, list]: """Extract outlines from binary image""" # Check if contours are being cut at image boundaries h, w = binary_image.shape # Add small border to prevent boundary cutting bordered_image = cv2.copyMakeBorder(binary_image, 5, 5, 5, 5, cv2.BORDER_CONSTANT, value=0) contours, _ = cv2.findContours( bordered_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE ) # Adjust contour coordinates back to original image space adjusted_contours = [] for contour in contours: adjusted_contour = contour - [5, 5] # Subtract border offset adjusted_contours.append(adjusted_contour) outline_image = np.full_like(binary_image, 255) return outline_image, adjusted_contours def round_edges(mask: np.ndarray, radius_mm: float, scaling_factor: float) -> np.ndarray: """Round mask edges using contour smoothing""" if radius_mm <= 0 or scaling_factor <= 0: return mask radius_px = max(1, int(radius_mm / scaling_factor)) if np.count_nonzero(mask) < 500: return cv2.dilate(cv2.erode(mask, np.ones((3,3))), np.ones((3,3))) contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) contours = [c for c in contours if cv2.contourArea(c) > 100] smoothed_contours = [] for contour in contours: try: resampled = resample_contour(contour, radius_px) resampled = resampled.astype(np.int32).reshape((-1, 1, 2)) smoothed_contours.append(resampled) except Exception as e: logger.warning(f"Error smoothing contour: {e}") smoothed_contours.append(contour) rounded = np.zeros_like(mask) cv2.drawContours(rounded, smoothed_contours, -1, 255, thickness=cv2.FILLED) return rounded def cleanup_memory(): """Clean up memory after processing""" if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect() logger.info("Memory cleanup completed") def cleanup_models(): """Unload models to free memory""" global paper_detector_global, u2net_global, birefnet if paper_detector_global is not None: del paper_detector_global paper_detector_global = None if u2net_global is not None: del u2net_global u2net_global = None if birefnet is not None: del birefnet birefnet = None cleanup_memory() def make_square(img: np.ndarray): """Make the image square by padding""" 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 predict_with_paper(image, paper_size, offset, offset_unit, finger_clearance=False): """Main prediction function using paper as reference""" logger.info(f"Starting prediction with image shape: {image.shape}") logger.info(f"Paper size: {paper_size}, Offset: {offset} {offset_unit}") # Convert offset to mm for internal calculations (DXF generation expects mm) if offset_unit == "inches": offset_mm = convert_units(offset, "inches", "mm") else: offset_mm = offset edge_radius = None if edge_radius is None or edge_radius == 0: edge_radius = 0.0001 if offset < 0: raise gr.Error("Offset Value Can't be negative") try: # Detect paper bounds and calculate scaling factor (always in mm for DXF) logger.info("Starting paper detection...") paper_contour, scaling_factor = detect_paper_bounds(image, paper_size, output_unit="mm") logger.info(f"Paper detected successfully with scaling factor: {scaling_factor:.6f} mm/px") except ReferenceBoxNotDetectedError as e: logger.error(f"Paper detection failed: {e}") return ( None, None, None, None, f"Error: {str(e)}" ) except Exception as e: logger.error(f"Unexpected error in paper detection: {e}") raise gr.Error(f"Error processing image: {str(e)}") try: # Get paper bounds with expansion rect = cv2.boundingRect(paper_contour) expansion = max(20, int(min(rect[2], rect[3]) * 0.05)) # 5% expansion x, y, w, h = rect x_min = max(0, x - expansion) y_min = max(0, y - expansion) x_max = min(image.shape[1], x + w + expansion) y_max = min(image.shape[0], y + h + expansion) # Process the expanded paper area cropped_image = image[y_min:y_max, x_min:x_max] crop_offset = (x_min, y_min) # Remove background objects_mask = remove_bg(cropped_image) # Resize mask back to cropped image size target_height = y_max - y_min target_width = x_max - x_min objects_mask_resized = cv2.resize(objects_mask, (target_width, target_height)) # Place back in full image space full_mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8) full_mask[y_min:y_max, x_min:x_max] = objects_mask_resized # Light filtering only - don't exclude paper area aggressively # Just remove small noise kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3)) objects_mask = cv2.morphologyEx(full_mask, cv2.MORPH_OPEN, kernel) # Debug: Save intermediate masks cv2.imwrite("./debug/objects_mask_after_processing.jpg", objects_mask) # Check if we actually have object pixels object_pixels = np.count_nonzero(objects_mask) if object_pixels < 300: # Minimum threshold raise NoObjectDetectedError("No significant object detected") # Validate single object validate_single_object(objects_mask, paper_contour) except (MultipleObjectsError, NoObjectDetectedError) as e: return ( None, None, None, None, f"Error: {str(e)}" ) except Exception as e: raise gr.Error(f"Error in object detection: {str(e)}") # Apply edge rounding if specified rounded_mask = objects_mask.copy() # Apply dilation for offset - more precise calculation using mm values if offset_mm > 0: offset_pixels = max(1, int(round(float(offset_mm) / scaling_factor))) if offset_pixels > 0: kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (offset_pixels*2+1, offset_pixels*2+1)) dilated_mask = cv2.dilate(rounded_mask, kernel, iterations=1) else: dilated_mask = rounded_mask.copy() else: dilated_mask = rounded_mask.copy() # Save original dilated mask for output Image.fromarray(dilated_mask).save("./outputs/scaled_mask_original.jpg") dilated_mask_orig = dilated_mask.copy() # Extract contours outlines, contours = extract_outlines(dilated_mask) try: # Generate DXF - scaling_factor should be in mm/px for proper DXF units dxf, finger_polygons, original_polygons = save_dxf_spline( contours, scaling_factor, # This should be mm/px image.shape[0], # Use original image height finger_clearance=(finger_clearance == "On") ) except FingerCutOverlapError as e: raise gr.Error(str(e)) # Create annotated image shrunked_img_contours = image.copy() if finger_clearance == "On": outlines = np.full_like(dilated_mask, 255) for poly in finger_polygons: try: coords = np.array([ (int(x / scaling_factor), int(image.shape[0] - y / scaling_factor)) for x, y in poly.exterior.coords ], np.int32).reshape((-1, 1, 2)) cv2.drawContours(shrunked_img_contours, [coords], -1, (0, 255, 0), thickness=2) cv2.drawContours(outlines, [coords], -1, 0, thickness=2) except Exception as e: logger.warning(f"Failed to draw finger cut: {e}") continue else: outlines = np.full_like(dilated_mask, 255) cv2.drawContours(shrunked_img_contours, contours, -1, (0, 255, 0), thickness=2) cv2.drawContours(outlines, contours, -1, 0, thickness=2) cleanup_models() # Format scaling info with proper unit display if offset_unit == "inches": offset_display = f"{offset} inches ({offset_mm:.3f} mm)" else: offset_display = f"{offset} mm" scale_info = f"Scale: {scaling_factor:.6f} mm/px | Paper: {paper_size} | Offset: {offset_display}" return ( shrunked_img_contours, outlines, dxf, dilated_mask_orig, scale_info ) def predict_full_paper(image, paper_size, offset_value_mm = 0.02,offset_unit='mm', enable_finger_cut='Off', selected_outputs=None): finger_flag = "On" if enable_finger_cut == "On" else "Off" # Always get all outputs from predict_with_paper ann, outlines, dxf_path, mask, scale_info = predict_with_paper( image, paper_size, offset=offset_value_mm, offset_unit= offset_unit, finger_clearance=finger_flag, ) # Return based on selected outputs return ( dxf_path, # Always return DXF ann if "Annotated Image" in selected_outputs else None, outlines if "Outlines" in selected_outputs else None, mask if "Mask" in selected_outputs else None, scale_info # Always return scaling info ) # Gradio Interface if __name__ == "__main__": os.makedirs("./outputs", exist_ok=True) with gr.Blocks(title="Paper-Based DXF Generator", theme=gr.themes.Soft()) as demo: # Example gallery gr.Markdown(""" # Paper-Based DXF Generator Upload an image with a single object placed on paper (A4, A3, or US Letter). The paper serves as a size reference for accurate DXF generation. """) with gr.Row(): gr.Markdown(""" **Instructions:** 1. Place a single object on paper 2. Select the correct paper size 3. Configure options as needed 4. Click Submit to generate DXF """) gr.Markdown(""" ### Tips for Best Results: - Ensure good lighting and clear paper edges - Place object completely at the center of the paper - Avoid shadows that might interfere with detection - Use high contrast between object and paper """) with gr.Row(): with gr.Column(): input_image = gr.Image( label="Input Image (Object on Paper)", type="numpy", height=400 ) paper_size = gr.Radio( choices=["A4", "A3", "US Letter"], value="A4", label="Paper Size", info="Select the paper size used in your image" ) # with gr.Group(): # gr.Markdown("### Contour Offset") # with gr.Row(): # offset_value_mm = gr.Number( # value=0.02, # label="Offset", # info="Expand contours outward by this amount", # precision=3, # minimum=0, # maximum=50 # ) # offset_unit = gr.Dropdown( # choices=["mm", "inches"], # value="mm", # label="Unit" # ) # with gr.Group(): # gr.Markdown("### Finger Cuts") # enable_finger_cut = gr.Radio( # choices=["On", "Off"], # value="Off", # label="Enable Finger Cuts", # info="Add circular cuts for easier handling" # ) output_options = gr.CheckboxGroup( choices=["Annotated Image", "Outlines", "Mask"], value=[], label="Additional Outputs", info="DXF is always included" ) submit_btn = gr.Button("Generate DXF", variant="primary", size="lg") with gr.Column(): with gr.Group(): gr.Markdown("### Generated Files") dxf_file = gr.File(label="DXF File", file_types=[".dxf"]) scale_info = gr.Textbox(label="Scaling Information", interactive=False) with gr.Group(): gr.Markdown("### Preview Images") output_image = gr.Image(label="Annotated Image", visible=False) outlines_image = gr.Image(label="Outlines", visible=False) mask_image = gr.Image(label="Mask", visible=False) def update_outputs_visibility(selected): return [ gr.update(visible="Annotated Image" in selected), gr.update(visible="Outlines" in selected), gr.update(visible="Mask" in selected) ] output_options.change( fn=update_outputs_visibility, inputs=output_options, outputs=[output_image, outlines_image, mask_image] ) submit_btn.click( fn=predict_full_paper, inputs=[ input_image, paper_size, gr.Number(value=0.02, visible=False), # Create hidden components gr.Textbox(value='mm', visible=False), gr.Textbox(value='Off', visible=False), output_options ], outputs=[dxf_file, output_image, outlines_image, mask_image, scale_info] ) demo.launch(share=True)