MLBench/Contours_Extraction
0
1 2 3import os4from pathlib import Path5from typing import List, Union6from PIL import Image7import ezdxf.units8import numpy as np9import torch10from torchvision import transforms11from ultralytics import YOLOWorld, YOLO12from ultralytics.engine.results import Results13from ultralytics.utils.plotting import save_one_box14from transformers import AutoModelForImageSegmentation15import cv216import ezdxf17import gradio as gr18import gc19from scalingtestupdated import calculate_scaling_factor20from scipy.interpolate import splprep, splev21from scipy.ndimage import gaussian_filter1d22import json23import time24import signal25from shapely.ops import unary_union26from shapely.geometry import MultiPolygon, GeometryCollection, Polygon, Point27from u2netp import U2NETP # Add U2NETP import28import logging29import shutil30 31# Initialize logging32logging.basicConfig(level=logging.INFO)33logger = logging.getLogger(__name__)34 35# Create cache directory for models36CACHE_DIR = os.path.join(os.path.dirname(__file__), ".cache")37os.makedirs(CACHE_DIR, exist_ok=True)38 39# Custom Exception Classes40class TimeoutReachedError(Exception):41 pass42 43class BoundaryOverlapError(Exception):44 pass45 46class TextOverlapError(Exception):47 pass48 49class ReferenceBoxNotDetectedError(Exception):50 """Raised when the Reference coin cannot be detected in the image"""51 pass52 53class FingerCutOverlapError(Exception):54 """Raised when finger cuts overlap with existing geometry"""55 def __init__(self, message="There was an overlap with fingercuts... Please try again to generate dxf."):56 super().__init__(message)57 58# ===== LAZY LOADING - REPLACE THE GLOBAL MODEL INITIALIZATION =====59# Instead of loading models at startup, declare them as None60print("Initializing lazy model loading...")61reference_detector_global = None62u2net_global = None63birefnet = None64 65# Model paths - use absolute paths for Docker66reference_model_path = os.path.join(CACHE_DIR, "best1.pt")67u2net_model_path = os.path.join(CACHE_DIR, "u2netp.pth")68 69# Copy model files to cache if they don't exist - with error handling70def ensure_model_files():71 if not os.path.exists(reference_model_path):72 if os.path.exists("best1.pt"):73 shutil.copy("best1.pt", reference_model_path)74 else:75 raise FileNotFoundError("best1.pt model file not found")76 if not os.path.exists(u2net_model_path):77 if os.path.exists("u2netp.pth"):78 shutil.copy("u2netp.pth", u2net_model_path)79 else:80 raise FileNotFoundError("u2netp.pth model file not found")81 82# Call this at startup83ensure_model_files()84 85# device = "cpu"86# torch.set_float32_matmul_precision(["high", "highest"][0])87 88# ===== LAZY LOADING FUNCTIONS - ADD THESE =====89def get_reference_detector():90 """Lazy load reference detector model"""91 global reference_detector_global92 if reference_detector_global is None:93 logger.info("Loading reference detector model...")94 reference_detector_global = YOLO(reference_model_path)95 logger.info("Reference detector loaded successfully")96 return reference_detector_global97 98def get_u2net():99 """Lazy load U2NETP model"""100 global u2net_global101 if u2net_global is None:102 logger.info("Loading U2NETP model...")103 u2net_global = U2NETP(3, 1)104 u2net_global.load_state_dict(torch.load(u2net_model_path, map_location="cpu"))105 u2net_global.to(device)106 u2net_global.eval()107 logger.info("U2NETP model loaded successfully")108 return u2net_global109def load_birefnet_model():110 """Load BiRefNet model from HuggingFace"""111 from transformers import AutoModelForImageSegmentation112 return AutoModelForImageSegmentation.from_pretrained(113 'ZhengPeng7/BiRefNet', 114 trust_remote_code=True115 )116def get_birefnet():117 """Lazy load BiRefNet model"""118 global birefnet119 if birefnet is None:120 logger.info("Loading BiRefNet model...")121 birefnet = load_birefnet_model()122 birefnet.to(device)123 birefnet.eval()124 logger.info("BiRefNet model loaded successfully")125 return birefnet126 127 128 129 130device = "cpu"131torch.set_float32_matmul_precision(["high", "highest"][0])132 133# Move models to device134# u2net_global.to(device)135# u2net_global.eval()136# birefnet.to(device)137# birefnet.eval()138 139# Define transforms140transform_image = transforms.Compose([141 transforms.Resize((1024, 1024)),142 transforms.ToTensor(),143 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),144])145 146def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:147 """Remove background using U2NETP model specifically for reference objects"""148 try:149 u2net_model = get_u2net() # <-- ADD THIS LINE150 151 image_pil = Image.fromarray(image)152 transform_u2netp = transforms.Compose([153 transforms.Resize((320, 320)),154 transforms.ToTensor(),155 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),156 ])157 158 input_tensor = transform_u2netp(image_pil).unsqueeze(0).to(device)159 160 with torch.no_grad():161 outputs = u2net_model(input_tensor) # <-- CHANGE FROM u2net_global162 163 pred = outputs[0]164 pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)165 pred_np = pred.squeeze().cpu().numpy()166 pred_np = cv2.resize(pred_np, (image_pil.width, image_pil.height))167 pred_np = (pred_np * 255).astype(np.uint8)168 169 return pred_np170 except Exception as e:171 logger.error(f"Error in U2NETP background removal: {e}")172 raise173 174def remove_bg(image: np.ndarray) -> np.ndarray:175 """Remove background using BiRefNet model for main objects"""176 try:177 birefnet_model = get_birefnet() # <-- ADD THIS LINE178 179 image = Image.fromarray(image)180 input_images = transform_image(image).unsqueeze(0).to(device)181 182 with torch.no_grad():183 preds = birefnet_model(input_images)[-1].sigmoid().cpu() # <-- CHANGE FROM birefnet184 pred = preds[0].squeeze()185 186 pred_pil: Image = transforms.ToPILImage()(pred)187 188 scale_ratio = 1024 / max(image.size)189 scaled_size = (int(image.size[0] * scale_ratio), int(image.size[1] * scale_ratio))190 191 return np.array(pred_pil.resize(scaled_size))192 except Exception as e:193 logger.error(f"Error in BiRefNet background removal: {e}")194 raise195 196def resize_img(img: np.ndarray, resize_dim):197 return np.array(Image.fromarray(img).resize(resize_dim))198 199def make_square(img: np.ndarray):200 """Make the image square by padding"""201 height, width = img.shape[:2]202 max_dim = max(height, width)203 204 pad_height = (max_dim - height) // 2205 pad_width = (max_dim - width) // 2206 207 pad_height_extra = max_dim - height - 2 * pad_height208 pad_width_extra = max_dim - width - 2 * pad_width209 210 if len(img.shape) == 3: # Color image211 padded = np.pad(212 img,213 (214 (pad_height, pad_height + pad_height_extra),215 (pad_width, pad_width + pad_width_extra),216 (0, 0),217 ),218 mode="edge",219 )220 else: # Grayscale image221 padded = np.pad(222 img,223 (224 (pad_height, pad_height + pad_height_extra),225 (pad_width, pad_width + pad_width_extra),226 ),227 mode="edge",228 )229 230 return padded231 232 233def detect_reference_square(img) -> tuple:234 """Detect reference square in the image and ignore other coins"""235 try:236 reference_detector = get_reference_detector() # <-- ADD THIS LINE237 238 res = reference_detector.predict(img, conf=0.70) # <-- CHANGE FROM reference_detector_global239 if not res or len(res) == 0 or len(res[0].boxes) == 0:240 raise ReferenceBoxNotDetectedError("Unable to detect the reference coin in the image.")241 242 # Get all detected boxes243 boxes = res[0].cpu().boxes.xyxy244 245 # Find the largest box (most likely the reference coin)246 largest_box = None247 max_area = 0248 for box in boxes:249 x_min, y_min, x_max, y_max = box250 area = (x_max - x_min) * (y_max - y_min)251 if area > max_area:252 max_area = area253 largest_box = box254 255 return (256 save_one_box(largest_box.unsqueeze(0), img, save=False),257 largest_box258 )259 except Exception as e:260 if not isinstance(e, ReferenceBoxNotDetectedError):261 logger.error(f"Error in reference square detection: {e}")262 raise ReferenceBoxNotDetectedError("Error detecting reference coin. Please try again with a clearer image.")263 raise264 265 266 267 268 269 270 271 272def exclude_scaling_box(273 image: np.ndarray,274 bbox: np.ndarray,275 orig_size: tuple,276 processed_size: tuple,277 expansion_factor: float = 1.2,278) -> np.ndarray:279 x_min, y_min, x_max, y_max = map(int, bbox)280 scale_x = processed_size[1] / orig_size[1]281 scale_y = processed_size[0] / orig_size[0]282 283 x_min = int(x_min * scale_x)284 x_max = int(x_max * scale_x)285 y_min = int(y_min * scale_y)286 y_max = int(y_max * scale_y)287 288 box_width = x_max - x_min289 box_height = y_max - y_min290 291 expanded_x_min = max(0, int(x_min - (expansion_factor - 1) * box_width / 2))292 expanded_x_max = min(293 image.shape[1], int(x_max + (expansion_factor - 1) * box_width / 2)294 )295 expanded_y_min = max(0, int(y_min - (expansion_factor - 1) * box_height / 2))296 expanded_y_max = min(297 image.shape[0], int(y_max + (expansion_factor - 1) * box_height / 2)298 )299 300 image[expanded_y_min:expanded_y_max, expanded_x_min:expanded_x_max] = 0301 return image302 303 304 305 306 307def resample_contour(contour, edge_radius_px: int = 0):308 """Resample contour with radius-aware smoothing and periodic handling."""309 logger.info(f"Starting resample_contour with contour of shape {contour.shape}")310 311 num_points = 1500312 sigma = max(2, int(edge_radius_px) // 4) # Adjust sigma based on radius313 314 if len(contour) < 4: # Need at least 4 points for spline with periodic condition315 error_msg = f"Contour must have at least 4 points, but has {len(contour)} points."316 logger.error(error_msg)317 raise ValueError(error_msg)318 319 try:320 contour = contour[:, 0, :]321 logger.debug(f"Reshaped contour to shape {contour.shape}")322 323 # Ensure contour is closed by making start and end points the same324 if not np.array_equal(contour[0], contour[-1]):325 contour = np.vstack([contour, contour[0]])326 327 # Create periodic spline representation328 tck, u = splprep(contour.T, u=None, s=0, per=True)329 330 # Evaluate spline at evenly spaced points331 u_new = np.linspace(u.min(), u.max(), num_points)332 x_new, y_new = splev(u_new, tck, der=0)333 334 # Apply Gaussian smoothing with wrap-around335 if sigma > 0:336 x_new = gaussian_filter1d(x_new, sigma=sigma, mode='wrap')337 y_new = gaussian_filter1d(y_new, sigma=sigma, mode='wrap')338 339 # Re-close the contour after smoothing340 x_new[-1] = x_new[0]341 y_new[-1] = y_new[0]342 343 result = np.array([x_new, y_new]).T344 logger.info(f"Completed resample_contour with result shape {result.shape}")345 return result346 347 except Exception as e:348 logger.error(f"Error in resample_contour: {e}")349 raise350 351 352 353 354 355 356# def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=False):357# doc = ezdxf.new(units=ezdxf.units.MM)358# doc.header["$INSUNITS"] = ezdxf.units.MM359# msp = doc.modelspace()360# final_polygons_inch = []361# finger_centers = []362# original_polygons = []363 364# for contour in inflated_contours:365# try:366# # Removed the second parameter since it was causing the error367# resampled_contour = resample_contour(contour) 368 369# points_inch = [(x * scaling_factor, (height - y) * scaling_factor) 370# for x, y in resampled_contour]371 372# if len(points_inch) < 3:373# continue374 375# tool_polygon = build_tool_polygon(points_inch)376# original_polygons.append(tool_polygon)377 378# if finger_clearance:379# try:380# tool_polygon, center = place_finger_cut_adjusted(381# tool_polygon, points_inch, finger_centers, final_polygons_inch382# )383# except FingerCutOverlapError:384# tool_polygon = original_polygons[-1]385 386# exterior_coords = polygon_to_exterior_coords(tool_polygon)387# if len(exterior_coords) < 3:388# continue389 390# msp.add_spline(exterior_coords, degree=3, dxfattribs={"layer": "TOOLS"})391# final_polygons_inch.append(tool_polygon)392 393# except ValueError as e:394# logger.warning(f"Skipping contour: {e}")395 396# dxf_filepath = os.path.join("./outputs", "out.dxf")397# doc.saveas(dxf_filepath)398# return dxf_filepath, final_polygons_inch, original_polygons399 400 401 402 403def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=False):404 doc = ezdxf.new(units=ezdxf.units.MM)405 doc.header["$INSUNITS"] = ezdxf.units.MM406 msp = doc.modelspace()407 final_polygons_inch = []408 finger_centers = []409 original_polygons = []410 411 # Scale correction factor based on your analysis412 scale_correction = 1.079413 414 for contour in inflated_contours:415 try:416 resampled_contour = resample_contour(contour) 417 418 points_inch = [(x * scaling_factor, (height - y) * scaling_factor) 419 for x, y in resampled_contour]420 421 if len(points_inch) < 3:422 continue423 424 tool_polygon = build_tool_polygon(points_inch)425 original_polygons.append(tool_polygon)426 427 if finger_clearance:428 try:429 tool_polygon, center = place_finger_cut_adjusted(430 tool_polygon, points_inch, finger_centers, final_polygons_inch431 )432 except FingerCutOverlapError:433 tool_polygon = original_polygons[-1]434 435 exterior_coords = polygon_to_exterior_coords(tool_polygon)436 if len(exterior_coords) < 3:437 continue438 439 # Apply scale correction AFTER finger cuts and polygon adjustments440 corrected_coords = [(x * scale_correction, y * scale_correction) for x, y in exterior_coords]441 442 msp.add_spline(corrected_coords, degree=3, dxfattribs={"layer": "TOOLS"})443 final_polygons_inch.append(tool_polygon)444 445 except ValueError as e:446 logger.warning(f"Skipping contour: {e}")447 448 dxf_filepath = os.path.join("./outputs", "out.dxf")449 doc.saveas(dxf_filepath)450 return dxf_filepath, final_polygons_inch, original_polygons451 452 453 454 455 456def build_tool_polygon(points_inch):457 return Polygon(points_inch)458 459 460 461def polygon_to_exterior_coords(poly):462 logger.info(f"Starting polygon_to_exterior_coords with input geometry type: {poly.geom_type}")463 464 try:465 # 1) If it's a GeometryCollection or MultiPolygon, fuse everything into one shape466 if poly.geom_type == "GeometryCollection" or poly.geom_type == "MultiPolygon":467 logger.debug(f"Performing unary_union on {poly.geom_type}")468 unified = unary_union(poly)469 if unified.is_empty:470 logger.warning("unary_union produced an empty geometry; returning empty list")471 return []472 # If union still yields multiple disjoint pieces, pick the largest Polygon473 if unified.geom_type == "GeometryCollection" or unified.geom_type == "MultiPolygon":474 largest = None475 max_area = 0.0476 for g in getattr(unified, "geoms", []):477 if hasattr(g, "area") and g.area > max_area and hasattr(g, "exterior"):478 max_area = g.area479 largest = g480 if largest is None:481 logger.warning("No valid Polygon found in unified geometry; returning empty list")482 return []483 poly = largest484 else:485 # Now unified should be a single Polygon or LinearRing486 poly = unified487 488 # 2) At this point, we must have a single Polygon (or something with an exterior)489 if not hasattr(poly, "exterior") or poly.exterior is None:490 logger.warning("Input geometry has no exterior ring; returning empty list")491 return []492 493 raw_coords = list(poly.exterior.coords)494 total = len(raw_coords)495 logger.info(f"Extracted {total} raw exterior coordinates")496 497 if total == 0:498 return []499 500 # 3) Subsample coordinates to at most 100 points (evenly spaced)501 max_pts = 100502 if total > max_pts:503 step = total // max_pts504 sampled = [raw_coords[i] for i in range(0, total, step)]505 # Ensure we include the last point to close the loop506 if sampled[-1] != raw_coords[-1]:507 sampled.append(raw_coords[-1])508 logger.info(f"Downsampled perimeter from {total} to {len(sampled)} points")509 return sampled510 else:511 return raw_coords512 513 except Exception as e:514 logger.error(f"Error in polygon_to_exterior_coords: {e}")515 return []516 517 518 519 520 521 522 523 524def place_finger_cut_adjusted(525 tool_polygon: Polygon,526 points_inch: list,527 existing_centers: list,528 all_polygons: list,529 circle_diameter: float = 25.4,530 min_gap: float = 0.5,531 max_attempts: int = 100532) -> (Polygon, tuple):533 logger.info(f"Starting place_finger_cut_adjusted with {len(points_inch)} input points")534 535 from shapely.geometry import Point536 import numpy as np537 import time538 import random539 540 # Fallback: if we run out of time or attempts, place in the "middle" of the outline541 def fallback_solution():542 logger.warning("Using fallback approach for finger cut placement")543 # Pick the midpoint of the original outline as a last-resort center544 fallback_center = points_inch[len(points_inch) // 2]545 r = circle_diameter / 2.0546 fallback_circle = Point(fallback_center).buffer(r, resolution=32)547 try:548 union_poly = tool_polygon.union(fallback_circle)549 except Exception as e:550 logger.warning(f"Fallback union failed ({e}); trying buffer-union fallback")551 union_poly = tool_polygon.buffer(0).union(fallback_circle.buffer(0))552 553 existing_centers.append(fallback_center)554 logger.info(f"Fallback finger cut placed at {fallback_center}")555 return union_poly, fallback_center556 557 # Precompute values558 r = circle_diameter / 2.0559 needed_center_dist = circle_diameter + min_gap560 561 # 1) Get perimeter coordinates of this polygon562 raw_perimeter = polygon_to_exterior_coords(tool_polygon)563 if not raw_perimeter:564 logger.warning("No valid exterior coords found; using fallback immediately")565 return fallback_solution()566 567 # 2) Possibly subsample to at most 100 perimeter points568 if len(raw_perimeter) > 100:569 step = len(raw_perimeter) // 100570 perimeter_coords = raw_perimeter[::step]571 logger.info(f"Subsampled perimeter from {len(raw_perimeter)} to {len(perimeter_coords)} points")572 else:573 perimeter_coords = raw_perimeter[:]574 575 # 3) Randomize the order to avoid bias576 indices = list(range(len(perimeter_coords)))577 random.shuffle(indices)578 logger.debug(f"Shuffled perimeter indices for candidate order")579 580 # 4) Non-blocking timeout setup581 start_time = time.time()582 timeout_secs = 5.0 # leave ~0.1s margin583 584 attempts = 0585 try:586 while attempts < max_attempts:587 # 5) Abort if we're running out of time588 if time.time() - start_time > timeout_secs - 0.1:589 logger.warning(f"Approaching timeout after {attempts} attempts")590 return fallback_solution()591 592 # 6) For each shuffled perimeter point, try small offsets593 for idx in indices:594 # Check timeout inside the loop as well595 if time.time() - start_time > timeout_secs - 0.05:596 logger.warning("Timeout during candidate-point loop")597 return fallback_solution()598 599 cx, cy = perimeter_coords[idx]600 # Try five small offsets: (0,0), (±min_gap/2, 0), (0, ±min_gap/2)601 for dx, dy in [(0, 0), (-min_gap/2, 0), (min_gap/2, 0), (0, -min_gap/2), (0, min_gap/2)]:602 candidate_center = (cx + dx, cy + dy)603 604 # 6a) Check distance to existing finger centers605 too_close_finger = any(606 np.hypot(candidate_center[0] - ex, candidate_center[1] - ey) 607 < needed_center_dist 608 for (ex, ey) in existing_centers609 )610 if too_close_finger:611 continue612 613 # 6b) Build candidate circle with reduced resolution for speed614 candidate_circle = Point(candidate_center).buffer(r, resolution=32)615 616 # 6c) Must overlap ≥30% with this polygon617 try:618 inter_area = tool_polygon.intersection(candidate_circle).area619 except Exception:620 continue621 622 if inter_area < 0.3 * candidate_circle.area:623 continue624 625 # 6d) Must not intersect or even "touch" any other polygon (buffered by min_gap)626 invalid = False627 for other_poly in all_polygons:628 if other_poly.equals(tool_polygon):629 # Don't compare against itself630 continue631 # Buffer the other polygon by min_gap to enforce a strict clearance632 if other_poly.buffer(min_gap).intersects(candidate_circle) or \633 other_poly.buffer(min_gap).touches(candidate_circle):634 invalid = True635 break636 if invalid:637 continue638 639 # 6e) Candidate passes all tests → union and return640 try:641 union_poly = tool_polygon.union(candidate_circle)642 # If union is a MultiPolygon (more than one piece), reject643 if union_poly.geom_type == "MultiPolygon" and len(union_poly.geoms) > 1:644 continue645 # If union didn't change anything (no real cut), reject646 if union_poly.equals(tool_polygon):647 continue648 except Exception:649 continue650 651 existing_centers.append(candidate_center)652 logger.info(f"Finger cut placed successfully at {candidate_center} after {attempts} attempts")653 return union_poly, candidate_center654 655 attempts += 1656 # If we've done half the attempts and we're near timeout, bail out657 if attempts >= (max_attempts // 2) and (time.time() - start_time) > timeout_secs * 0.8:658 logger.warning(f"Approaching timeout (attempt {attempts})")659 return fallback_solution()660 661 logger.debug(f"Completed iteration {attempts}/{max_attempts}")662 663 # If we exit loop without finding a valid spot664 logger.warning(f"No valid spot after {max_attempts} attempts, using fallback")665 return fallback_solution()666 667 except Exception as e:668 logger.error(f"Error in place_finger_cut_adjusted: {e}")669 return fallback_solution()670 671 672 673 674 675 676 677 678 679 680def extract_outlines(binary_image: np.ndarray) -> tuple:681 contours, _ = cv2.findContours(682 binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE683 )684 685 outline_image = np.full_like(binary_image, 255) # White background686 687 return outline_image, contours688 689 690 691 692def round_edges(mask: np.ndarray, radius_mm: float, scaling_factor: float) -> np.ndarray:693 """Rounds mask edges using contour smoothing."""694 if radius_mm <= 0 or scaling_factor <= 0:695 return mask696 697 radius_px = max(1, int(radius_mm / scaling_factor)) # Ensure min 1px698 699 # Handle small objects700 if np.count_nonzero(mask) < 500: # Small object threshold701 return cv2.dilate(cv2.erode(mask, np.ones((3,3))), np.ones((3,3)))702 703 # Existing contour processing with improvements:704 contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)705 706 # NEW: Filter small contours707 contours = [c for c in contours if cv2.contourArea(c) > 100]708 smoothed_contours = []709 710 for contour in contours:711 try:712 # Resample with radius-based smoothing713 resampled = resample_contour(contour, radius_px)714 resampled = resampled.astype(np.int32).reshape((-1, 1, 2))715 smoothed_contours.append(resampled)716 except Exception as e:717 logger.warning(f"Error smoothing contour: {e}")718 smoothed_contours.append(contour) # Fallback to original contour719 720 # Draw smoothed contours721 rounded = np.zeros_like(mask)722 cv2.drawContours(rounded, smoothed_contours, -1, 255, thickness=cv2.FILLED)723 724 return rounded725 726def cleanup_memory():727 """Clean up memory after processing"""728 if torch.cuda.is_available():729 torch.cuda.empty_cache()730 gc.collect()731 logger.info("Memory cleanup completed")732 733def cleanup_models():734 """Unload models to free memory"""735 global reference_detector_global, u2net_global, birefnet736 if reference_detector_global is not None:737 del reference_detector_global738 reference_detector_global = None739 if u2net_global is not None:740 del u2net_global741 u2net_global = None742 if birefnet is not None:743 del birefnet744 birefnet = None745 cleanup_memory()746 747def predict_og(image, offset, offset_unit, edge_radius, finger_clearance=False):748 coin_size_mm = 20.0749 750 if offset_unit == "inches":751 offset *= 25.4752 753 if edge_radius is None or edge_radius == 0:754 edge_radius = 0.0001755 756 if offset < 0:757 raise gr.Error("Offset Value Can't be negative")758 759 try:760 reference_obj_img, scaling_box_coords = detect_reference_square(image)761 except ReferenceBoxNotDetectedError as e:762 return (763 None,764 None,765 None,766 None,767 f"Error: {str(e)}"768 )769 except Exception as e:770 raise gr.Error(f"Error processing image: {str(e)}")771 772 reference_obj_img = make_square(reference_obj_img)773 774 # Use U2NETP for reference object background removal775 reference_square_mask = remove_bg_u2netp(reference_obj_img)776 reference_square_mask = resize_img(reference_square_mask, reference_obj_img.shape[:2][::-1])777 778 try:779 scaling_factor = calculate_scaling_factor(780 target_image=reference_square_mask,781 reference_obj_size_mm=coin_size_mm,782 feature_detector="ORB",783 )784 except Exception as e:785 scaling_factor = None786 logger.warning(f"Error calculating scaling factor: {e}")787 788 if not scaling_factor:789 ref_size_px = (reference_square_mask.shape[0] + reference_square_mask.shape[1]) / 2790 scaling_factor = 20.0 / ref_size_px791 logger.info(f"Fallback scaling: {scaling_factor:.4f} mm/px using 20mm reference")792 793 # Use BiRefNet for main object background removal794 orig_size = image.shape[:2]795 objects_mask = remove_bg(image)796 processed_size = objects_mask.shape[:2]797 798 # REMOVE ALL COINS from mask:799 # res = reference_detector_global.predict(image, conf=0.05)800 res = get_reference_detector().predict(image, conf=0.05)801 boxes = res[0].cpu().boxes.xyxy if res and len(res) > 0 else []802 803 for box in boxes:804 objects_mask = exclude_scaling_box(805 objects_mask,806 box,807 orig_size,808 processed_size,809 expansion_factor=1.2,810 )811 812 objects_mask = resize_img(objects_mask, (image.shape[1], image.shape[0]))813 814 offset_pixels = (float(offset) / scaling_factor) * 2 + 1 if scaling_factor else 1815 dilated_mask = cv2.dilate(objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8))816 Image.fromarray(dilated_mask).save("./outputs/scaled_mask_original.jpg")817 dilated_mask_orig = dilated_mask.copy()818 819 #if edge_radius > 0:820 # Use morphological rounding instead of contour-based821 rounded_mask = round_edges(objects_mask, edge_radius, scaling_factor)822 #else:823 #rounded_mask = objects_mask.copy()824 825 # Apply dilation AFTER rounding826 offset_pixels = (float(offset) / scaling_factor) * 2 + 1 if scaling_factor else 1827 kernel = np.ones((int(offset_pixels), int(offset_pixels)), np.uint8)828 dilated_mask = cv2.dilate(rounded_mask, kernel)829 830 831 832 outlines, contours = extract_outlines(dilated_mask)833 834 try:835 dxf, finger_polygons, original_polygons = save_dxf_spline(836 contours,837 scaling_factor,838 processed_size[0],839 finger_clearance=(finger_clearance == "On")840 )841 except FingerCutOverlapError as e:842 raise gr.Error(str(e))843 844 shrunked_img_contours = image.copy()845 846 if finger_clearance == "On":847 outlines = np.full_like(dilated_mask, 255)848 for poly in finger_polygons:849 try:850 coords = np.array([851 (int(x / scaling_factor), int(processed_size[0] - y / scaling_factor))852 for x, y in poly.exterior.coords853 ], np.int32).reshape((-1, 1, 2))854 855 cv2.drawContours(shrunked_img_contours, [coords], -1, 0, thickness=2)856 cv2.drawContours(outlines, [coords], -1, 0, thickness=2)857 except Exception as e:858 logger.warning(f"Failed to draw finger cut: {e}")859 continue860 else:861 outlines = np.full_like(dilated_mask, 255)862 cv2.drawContours(shrunked_img_contours, contours, -1, 0, thickness=2)863 cv2.drawContours(outlines, contours, -1, 0, thickness=2)864 cleanup_models()865 866 return (867 shrunked_img_contours,868 outlines,869 dxf,870 dilated_mask_orig,871 f"{scaling_factor:.4f}")872 873 874def predict_simple(image):875 """876 Only image in → returns (annotated, outlines, dxf, mask).877 Uses offset=0 mm, no fillet, no finger-cut.878 """879 ann, outlines, dxf_path, mask, _ = predict_og(880 image,881 offset=0,882 offset_unit="mm",883 edge_radius=0,884 finger_clearance="Off",885 )886 return ann, outlines, dxf_path, mask887 888def predict_middle(image, enable_fillet, fillet_value_mm):889 """890 image + (On/Off) fillet toggle + fillet radius → returns (annotated, outlines, dxf, mask).891 Uses offset=0 mm, finger-cut off.892 """893 radius = fillet_value_mm if enable_fillet == "On" else 0894 ann, outlines, dxf_path, mask, _ = predict_og(895 image,896 offset=0,897 offset_unit="mm",898 edge_radius=radius,899 finger_clearance="Off",900 )901 return ann, outlines, dxf_path, mask902 903def predict_full(image, enable_fillet, fillet_value_mm, enable_finger_cut, selected_outputs):904 """905 Returns DXF + conditionally selected additional outputs906 Always returns exactly 4 values to match output components907 """908 radius = fillet_value_mm if enable_fillet == "On" else 0909 finger_flag = "On" if enable_finger_cut == "On" else "Off"910 911 # Always get all outputs from predict_og912 ann, outlines, dxf_path, mask, _ = predict_og(913 image,914 offset=0,915 offset_unit="mm",916 edge_radius=radius,917 finger_clearance=finger_flag,918 )919 920 # Always return 4 values to match the 4 output components921 return (922 dxf_path, # Always return DXF923 ann if "Annotated Image" in selected_outputs else None,924 outlines if "Outlines" in selected_outputs else None,925 mask if "Mask" in selected_outputs else None926 )927 928 929 930if __name__ == "__main__":931 os.makedirs("./outputs", exist_ok=True)932 933 with gr.Blocks() as demo:934 input_image = gr.Image(label="Input Image", type="numpy")935 936 enable_fillet = gr.Radio(937 choices=["On", "Off"],938 value="Off",939 label="Enable Fillet",940 interactive=True941 )942 943 fillet_value_mm = gr.Slider(944 minimum=0,945 maximum=20,946 step=1,947 value=5,948 label="Edge Radius (mm)",949 visible=False,950 interactive=True951 )952 953 enable_finger_cut = gr.Radio(954 choices=["On", "Off"],955 value="Off",956 label="Enable Finger Cut"957 )958 output_options = gr.CheckboxGroup(959 choices=["Annotated Image", "Outlines", "Mask"],960 value=[],961 label="Additional Outputs (DXF is always included)"962 )963 def toggle_fillet(choice):964 if choice == "On":965 return gr.update(visible=True)966 return gr.update(visible=False, value=0)967 968 enable_fillet.change(969 fn=toggle_fillet,970 inputs=enable_fillet,971 outputs=fillet_value_mm972 )973 974 dxf_file = gr.File(label="DXF file")975 output_image = gr.Image(label="Annotated Image", visible=False)976 outlines_image = gr.Image(label="Outlines", visible=False) 977 mask_image = gr.Image(label="Mask", visible=False)978 979 submit_btn = gr.Button("Submit")980 981 # Function to update output visibility982 def update_outputs_visibility(selected):983 return [984 gr.update(visible="Annotated Image" in selected),985 gr.update(visible="Outlines" in selected),986 gr.update(visible="Mask" in selected)987 ]988 989 # Connect visibility updates990 output_options.change(991 fn=update_outputs_visibility,992 inputs=output_options,993 outputs=[output_image, outlines_image, mask_image]994 )995 996 # Dynamic output list for submit button997 def get_outputs(selected):998 outputs = [dxf_file]999 if "Annotated Image" in selected:1000 outputs.append(output_image)1001 if "Outlines" in selected:1002 outputs.append(outlines_image) 1003 if "Mask" in selected:1004 outputs.append(mask_image)1005 return outputs1006 1007 submit_btn.click(1008 fn=predict_full,1009 inputs=[input_image, enable_fillet, fillet_value_mm, enable_finger_cut, output_options],1010 outputs=[dxf_file, output_image, outlines_image, mask_image]1011 )1012 1013 # output_image = gr.Image(label="Output Image")1014 # outlines = gr.Image(label="Outlines of Objects")1015 # dxf_file = gr.File(label="DXF file")1016 # mask = gr.Image(label="Mask")1017 1018 # submit_btn = gr.Button("Submit")1019 1020 1021 # submit_btn.click(1022 # fn=predict_full,1023 # inputs=[input_image, enable_fillet, fillet_value_mm, enable_finger_cut],1024 # outputs=[output_image, outlines, dxf_file, mask]1025 # )1026 1027 demo.launch(share=True)1028 1029# import os1030# from pathlib import Path1031# from typing import List, Union1032# from PIL import Image1033# import ezdxf.units1034# import numpy as np1035# import torch1036# from torchvision import transforms1037# from ultralytics import YOLOWorld, YOLO1038# from ultralytics.engine.results import Results1039# from ultralytics.utils.plotting import save_one_box1040# from transformers import AutoModelForImageSegmentation1041# import cv21042# import ezdxf1043# import gradio as gr1044# import gc1045# from scalingtestupdated import calculate_scaling_factor1046# from scipy.interpolate import splprep, splev1047# from scipy.ndimage import gaussian_filter1d1048# import json1049# import time1050# import signal1051# from shapely.ops import unary_union1052# from shapely.geometry import MultiPolygon, GeometryCollection, Polygon, Point1053# from u2netp import U2NETP # Add U2NETP import1054# import logging1055# import shutil1056 1057# # Initialize logging1058# logging.basicConfig(level=logging.INFO)1059# logger = logging.getLogger(__name__)1060 1061# # Create cache directory for models1062# CACHE_DIR = os.path.join(os.path.dirname(__file__), ".cache")1063# os.makedirs(CACHE_DIR, exist_ok=True)1064 1065# # Custom Exception Classes1066# class TimeoutReachedError(Exception):1067# pass1068 1069# class BoundaryOverlapError(Exception):1070# pass1071 1072# class TextOverlapError(Exception):1073# pass1074 1075# class ReferenceBoxNotDetectedError(Exception):1076# """Raised when the Reference coin cannot be detected in the image"""1077# pass1078 1079# class FingerCutOverlapError(Exception):1080# """Raised when finger cuts overlap with existing geometry"""1081# def __init__(self, message="There was an overlap with fingercuts... Please try again to generate dxf."):1082# super().__init__(message)1083 1084# # ===== LAZY LOADING - REPLACE THE GLOBAL MODEL INITIALIZATION =====1085# # Instead of loading models at startup, declare them as None1086# print("Initializing lazy model loading...")1087# reference_detector_global = None1088# u2net_global = None1089# birefnet = None1090 1091# # Model paths - use absolute paths for Docker1092# reference_model_path = os.path.join(CACHE_DIR, "best1.pt")1093# u2net_model_path = os.path.join(CACHE_DIR, "u2netp.pth")1094 1095# # Copy model files to cache if they don't exist - with error handling1096# def ensure_model_files():1097# if not os.path.exists(reference_model_path):1098# if os.path.exists("best1.pt"):1099# shutil.copy("best1.pt", reference_model_path)1100# else:1101# raise FileNotFoundError("best1.pt model file not found")1102# if not os.path.exists(u2net_model_path):1103# if os.path.exists("u2netp.pth"):1104# shutil.copy("u2netp.pth", u2net_model_path)1105# else:1106# raise FileNotFoundError("u2netp.pth model file not found")1107 1108# # Call this at startup1109# ensure_model_files()1110 1111# # device = "cpu"1112# # torch.set_float32_matmul_precision(["high", "highest"][0])1113 1114# # ===== LAZY LOADING FUNCTIONS - ADD THESE =====1115# def get_reference_detector():1116# """Lazy load reference detector model"""1117# global reference_detector_global1118# if reference_detector_global is None:1119# logger.info("Loading reference detector model...")1120# reference_detector_global = YOLO(reference_model_path)1121# logger.info("Reference detector loaded successfully")1122# return reference_detector_global1123 1124# def get_u2net():1125# """Lazy load U2NETP model"""1126# global u2net_global1127# if u2net_global is None:1128# logger.info("Loading U2NETP model...")1129# u2net_global = U2NETP(3, 1)1130# u2net_global.load_state_dict(torch.load(u2net_model_path, map_location="cpu"))1131# u2net_global.to(device)1132# u2net_global.eval()1133# logger.info("U2NETP model loaded successfully")1134# return u2net_global1135# def load_birefnet_model():1136# """Load BiRefNet model from HuggingFace"""1137# from transformers import AutoModelForImageSegmentation1138# return AutoModelForImageSegmentation.from_pretrained(1139# 'ZhengPeng7/BiRefNet', 1140# trust_remote_code=True1141# )1142# def get_birefnet():1143# """Lazy load BiRefNet model"""1144# global birefnet1145# if birefnet is None:1146# logger.info("Loading BiRefNet model...")1147# birefnet = load_birefnet_model()1148# birefnet.to(device)1149# birefnet.eval()1150# logger.info("BiRefNet model loaded successfully")1151# return birefnet1152 1153 1154 1155 1156# device = "cpu"1157# torch.set_float32_matmul_precision(["high", "highest"][0])1158 1159# # Move models to device1160# # u2net_global.to(device)1161# # u2net_global.eval()1162# # birefnet.to(device)1163# # birefnet.eval()1164 1165# # Define transforms1166# transform_image = transforms.Compose([1167# transforms.Resize((1024, 1024)),1168# transforms.ToTensor(),1169# transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),1170# ])1171 1172# def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:1173# """Remove background using U2NETP model specifically for reference objects"""1174# try:1175# u2net_model = get_u2net() # <-- ADD THIS LINE1176 1177# image_pil = Image.fromarray(image)1178# transform_u2netp = transforms.Compose([1179# transforms.Resize((320, 320)),1180# transforms.ToTensor(),1181# transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),1182# ])1183 1184# input_tensor = transform_u2netp(image_pil).unsqueeze(0).to(device)1185 1186# with torch.no_grad():1187# outputs = u2net_model(input_tensor) # <-- CHANGE FROM u2net_global1188 1189# pred = outputs[0]1190# pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)1191# pred_np = pred.squeeze().cpu().numpy()1192# pred_np = cv2.resize(pred_np, (image_pil.width, image_pil.height))1193# pred_np = (pred_np * 255).astype(np.uint8)1194 1195# return pred_np1196# except Exception as e:1197# logger.error(f"Error in U2NETP background removal: {e}")1198# raise1199 1200# def remove_bg(image: np.ndarray) -> np.ndarray: