Upload 11 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+Test20.jpg filter=lfs diff=lfs merge=lfs -text
+Test21.jpg filter=lfs diff=lfs merge=lfs -text

app.py

@@ -0,0 +1,912 @@
+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
+from ezdxf.addons.text2path import make_paths_from_str
+from ezdxf import path
+from ezdxf.addons import text2path
+from ezdxf.enums import TextEntityAlignment
+from ezdxf.fonts.fonts import FontFace, get_font_face
+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 coin cannot be detected in the image"""
+    pass
+
+class BoundaryOverlapError(Exception):
+    """Raised when the optional boundary dimensions are too small and overlap with the inner contours."""
+    pass
+
+class TextOverlapError(Exception):
+    """Raised when the text overlaps with the inner contours (with a margin of 0.75)."""
+    pass
+class boundary_issue(Exception):
+    """Raised when bounds are given but rectangular boundary is no."""
+# ---------------------
+# 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, "coin_det.pt")
+if not os.path.exists(reference_model_path):
+    print("Caching YOLO reference model to", reference_model_path)
+    shutil.copy("coin_det.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, "coin_det.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.3)
+    if not res or len(res) == 0 or len(res[0].boxes) == 0:
+        raise ReferenceBoxNotDetectedError("Reference Coin 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): # works fine
+    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_adjusted(tool_polygon, points_inch, existing_centers, all_polygons, circle_diameter=1, min_gap=1, max_attempts=500): #1st best
+#     needed_center_distance = circle_diameter + min_gap
+#     radius = circle_diameter / 2.0
+#     import random
+#     for _ in range(max_attempts):
+#         idx = random.randint(0, len(points_inch) - 1)
+#         cx, cy = points_inch[idx]
+
+#         # Check if this point is too close to an existing center
+#         too_close = any(np.hypot(cx - ex_x, cy - ex_y) < needed_center_distance for ex_x, ex_y in existing_centers)
+#         if too_close:
+#             continue
+
+#         # Create the finger cut circle and try adding it to the tool
+#         circle_poly = Point((cx, cy)).buffer(radius, resolution=64)
+#         union_poly = tool_polygon.union(circle_poly)
+
+#         # Check for overlap and spacing with other tools
+#         overlap_with_others = False
+#         too_close_to_others = False
+
+#         for poly in all_polygons:
+#             if poly.equals(tool_polygon):
+#                 continue  # Skip comparing the tool to itself
+
+#             if union_poly.buffer(min_gap).intersects(poly) > 1e-6:
+#                 overlap_with_others = True
+#                 break
+
+#             if circle_poly.buffer(min_gap).intersects(poly) > 1e-6:
+#                 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
+
+import numpy as np
+from shapely.geometry import Point
+
+
+# def place_finger_cut_adjusted(tool_polygon, points_inch, existing_centers, all_polygons, circle_diameter=1.0, min_gap=0.35, max_attempts=2000): #Best best
+#     import random
+#     import numpy as np
+#     from shapely.geometry import Point
+    
+#     needed_center_distance = circle_diameter + min_gap
+#     radius = circle_diameter / 2.0
+#     attempts = 0
+#     indices = list(range(len(points_inch)))
+#     random.shuffle(indices)  # Shuffle indices for randomness
+    
+#     # Try a grid of adjustments around each candidate point
+#     adjustments = list(np.linspace(-0.15, 0.10, 7))  # More adjustment options
+    
+#     for i in indices:
+#         if attempts >= max_attempts:
+#             break
+            
+#         cx, cy = points_inch[i]
+        
+#         # Try small adjustments around the chosen candidate
+#         for dx in adjustments:
+#             for dy in adjustments:
+#                 attempts += 1
+#                 if attempts >= max_attempts:
+#                     break
+                    
+#                 candidate_center = (cx + dx, cy + dy)
+                
+#                 # Check distance from already placed centers
+#                 too_close_to_existing = False
+#                 for ex, ey in existing_centers:
+#                     if np.hypot(candidate_center[0] - ex, candidate_center[1] - ey) < needed_center_distance:
+#                         too_close_to_existing = True
+#                         break
+                
+#                 if too_close_to_existing:
+#                     continue
+                    
+#                 # Create circle polygon for this candidate
+#                 circle_poly = Point(candidate_center).buffer(radius, resolution=64)
+                
+#                 # Create the union with the tool polygon
+#                 union_poly = tool_polygon.union(circle_poly)
+                
+#                 # Buffer the circle to check minimum gap requirements
+#                 circle_buffer = circle_poly.buffer(min_gap, resolution=32)
+#                 coords = polygon_to_exterior_coords(union_poly)
+                
+#                 # Check against all other polygons for overlap or proximity issues
+#                 overlap = False
+#                 for poly in all_polygons:
+#                     if poly == tool_polygon:
+#                         continue  # Skip comparing to self
+#                     if len(coords) < 4:
+#                         # It's degenerate or not a valid polygon for your purposes; skip
+#                         break
+                            
+#                     # Check if the union overlaps with any other polygon
+#                     if union_poly.intersects(poly):
+#                         overlap = True
+#                         break
+                        
+#                     # Check if the buffered circle (circle + min_gap) intersects with any other polygon
+#                     if circle_buffer.intersects(poly):
+#                         overlap = True
+#                         break
+                
+#                 if not overlap:
+#                     # If candidate passes all checks, accept it
+#                     existing_centers.append(candidate_center)
+#                     return union_poly, candidate_center
+    
+#     print(f"Warning: Could not place a finger cut circle after {attempts} attempts. Consider adjusting parameters.")
+#     return None, None
+
+
+def place_finger_cut_adjusted(tool_polygon, points_inch, existing_centers, all_polygons, circle_diameter=1.0, min_gap=0.25, max_attempts=100):
+    import random
+    needed_center_distance = circle_diameter + min_gap
+    radius = circle_diameter / 2.0
+    attempts = 0
+    indices = list(range(len(points_inch)))
+    random.shuffle(indices)  # Shuffle indices for randomness
+
+    for i in indices:
+        if attempts >= max_attempts:
+            break
+        cx, cy = points_inch[i]
+        # Try small adjustments around the chosen candidate
+        for dx in np.linspace(-0.1, 0.1, 10):
+            for dy in np.linspace(-0.1, 0.1, 10):
+                candidate_center = (cx + dx, cy + dy)
+                # Check distance from already placed centers
+                if any(np.hypot(candidate_center[0] - ex, candidate_center[1] - ey) < needed_center_distance for ex, ey in existing_centers):
+                    continue
+                
+                union_poly= union_tool_and_circle(tool_polygon,candidate_center)
+                overlap = False
+                # Check against other tool polygons for overlap or proximity issues
+                for poly in all_polygons:
+                    if poly == tool_polygon:
+                        continue  
+                    if union_poly.intersects(poly) or union_poly.buffer(min_gap).intersects(poly):
+                        overlap = True
+                        break
+                if overlap:
+                    continue
+                # If candidate passes, accept it
+                existing_centers.append(candidate_center)
+                return union_poly, candidate_center
+        attempts += 1
+    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): # works
+    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_adjusted(tool_polygon, points_inch, finger_cut_centers, final_polygons_inch)
+                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, offset_unit, annotation_text="", image_height_in=None, image_width_in=None):
+    msp = doc.modelspace()
+    # Convert from mm if necessary
+    if offset_unit.lower() == "mm":
+        if boundary_length < 50:
+            boundary_length = boundary_length * 25.4
+        if boundary_width < 50:
+            boundary_width = boundary_width * 25.4
+        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
+
+    # Compute bounding box of inner contours
+    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
+
+    # Compute inner bounding box dimensions
+    inner_width = max_x - min_x
+    inner_length = max_y - min_y
+
+    # Set clearance margins
+    clearance_side = 0.25  # left/right clearance
+    clearance_tb = 0.25    # top/bottom clearance
+    if annotation_text.strip():
+        clearance_tb = 0.75
+
+    # Calculate center of inner contours
+    center_x = (min_x + max_x) / 2
+    center_y = (min_y + max_y) / 2
+
+    # Draw rectangle centered at (center_x, center_y)
+    left = center_x - boundary_width_in / 2
+    right = center_x + boundary_width_in / 2
+    bottom = center_y - boundary_length_in / 2
+    top = center_y + boundary_length_in / 2
+
+    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"})
+    
+    text_top = boundary_polygon.bounds[1] + 1
+    too_small = boundary_width_in < inner_width + 2 * clearance_side or boundary_length_in < inner_length + 2 * clearance_tb
+    if too_small:
+        raise BoundaryOverlapError("Error: The specified boundary dimensions are too small and overlap with the inner contours. Please provide larger values.")
+    if annotation_text.strip() and text_top > min_y - 0.75:
+        raise TextOverlapError("Error: The text is too close to the inner contours. Please increase boundary length.")
+    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_value: float,
+    offset_unit: str,         # "mm" or "inches"
+    finger_clearance: str,    # "Yes" or "No"
+    add_boundary: str,        # "Yes" or "No"
+    boundary_length: float,
+    boundary_width: float,
+    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 brightness and sharpness enhancement
+    if isinstance(image, np.ndarray):
+        pil_image = Image.fromarray(image)
+        enhanced_image = ImageEnhance.Sharpness(pil_image).enhance(1.5)
+        image = np.array(enhanced_image)
+
+    # ---------------------
+    # 1) Detect the drawer with YOLOWorld (or use original image if not detected)
+    # ---------------------
+    drawer_detected = True
+    try:
+        t = time.time()
+        drawer_img = yolo_detect(image)
+        print("Drawer detection completed in {:.2f} seconds".format(time.time() - t))
+    except DrawerNotDetectedError as e:
+        print(f"Drawer not detected: {e}, using original image.")
+        drawer_detected = False
+        drawer_img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+
+    # Process the image (either cropped drawer or original)
+    t = time.time()
+    if drawer_detected:
+        # For detected drawers: shrink and square
+        shrunked_img = make_square(shrink_bbox(drawer_img, 0.90))
+    else:
+        # For non-drawer images: keep original dimensions
+        shrunked_img = drawer_img  # Already in BGR format from above
+    del drawer_img
+    gc.collect()
+    print("Image processing completed in {:.2f} seconds".format(time.time() - t))
+
+    # ---------------------
+    # 2) Detect the reference box with YOLO (now works on either cropped or original image)
+    # ---------------------
+    try:
+        t = time.time()
+        reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
+        print("Reference coin detection completed in {:.2f} seconds".format(time.time() - t))
+    except ReferenceBoxNotDetectedError as e:
+        return None, None, None, None, f"Error: {str(e)}"
+
+    # ---------------------
+    # 3) Remove background of the reference box to compute scaling factor
+    # ---------------------
+    t = time.time()
+    reference_obj_img = make_square(reference_obj_img)
+    reference_square_mask = remove_bg_u2netp(reference_obj_img)
+    reference_square_mask= resize_img(reference_square_mask,(reference_obj_img.shape[1],reference_obj_img.shape[0]))
+    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(
+            target_image=reference_square_mask,
+            reference_obj_size_mm=0.955,
+            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 = 0.7
+        print("Using default scaling factor of 0.7 due to calculation error")
+    gc.collect()
+    print("Scaling factor determined: {}".format(scaling_factor))
+
+    # ---------------------
+    # 4) Optional boundary dimension checks (now without size limits)
+    # ---------------------
+    if add_boundary.lower() == "yes":
+        if offset_unit.lower() == "mm":
+            if boundary_length < 50:
+                boundary_length = boundary_length * 25.4
+            if boundary_width < 50:
+                boundary_width = boundary_width * 25.4
+            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
+        
+    # ---------------------
+    # 5) Remove background from the shrunked drawer image (main objects)
+    # ---------------------
+    if offset_unit.lower() == "mm":
+        if offset_value < 1:
+            offset_value = offset_value * 25.4
+        offset_inches = offset_value / 25.4
+    else:
+        offset_inches = offset_value
+
+    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))
+
+    # Dilate mask by offset_pixels
+    t = time.time()
+    offset_pixels = (offset_inches / scaling_factor) * 2 + 1 if scaling_factor != 0 else 1
+    dilated_mask = cv2.dilate(objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8))
+    del objects_mask
+    gc.collect()
+    print("Mask dilation completed in {:.2f} seconds".format(time.time() - t))
+
+    Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
+
+    # ---------------------
+    # 6) Extract outlines from the mask and convert them to DXF splines
+    # ---------------------
+    t = time.time()
+    outlines, contours = extract_outlines(dilated_mask)
+    print("Outline extraction completed in {:.2f} seconds".format(time.time() - t))
+
+    output_img = shrunked_img.copy()
+    del shrunked_img
+    gc.collect()
+
+    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))
+
+    # ---------------------
+    # Compute bounding box of inner tool contours BEFORE adding optional boundary
+    # ---------------------
+    inner_min_x = float("inf")
+    inner_min_y = float("inf")
+    inner_max_x = -float("inf")
+    inner_max_y = -float("inf")
+    for poly in final_polygons_inch:
+        b = poly.bounds
+        inner_min_x = min(inner_min_x, b[0])
+        inner_min_y = min(inner_min_y, b[1])
+        inner_max_x = max(inner_max_x, b[2])
+        inner_max_y = max(inner_max_y, b[3])
+
+    # ---------------------
+    # 7) Add optional rectangular boundary
+    # ---------------------
+    boundary_polygon = None
+    if add_boundary.lower() == "yes":
+        boundary_polygon = add_rectangular_boundary(
+            doc,
+            final_polygons_inch,
+            boundary_length,
+            boundary_width,
+            offset_unit,
+            annotation_text,
+            image_height_in=output_img.shape[0] * scaling_factor,
+            image_width_in=output_img.shape[1] * scaling_factor
+        )
+        if boundary_polygon is not None:
+            final_polygons_inch.append(boundary_polygon)
+        # else:
+        #     raise boundary_issue("Raised when bounds are given but rectangular boundary is no.")
+    # ---------------------
+    # 8) Add annotation text (if provided) in the DXF
+    # ---------------------
+    msp = doc.modelspace()
+    
+    if annotation_text.strip():
+        if boundary_polygon is not None:
+            text_x = ((inner_min_x + inner_max_x) / 2.0) - (int(len(annotation_text.strip()) / 2.0))
+            text_height_dxf = 0.75  
+            text_y_dxf = boundary_polygon.bounds[1] + 0.25
+            font = get_font_face("Arial")
+            paths = text2path.make_paths_from_str(
+                annotation_text.strip().upper(),
+                font=font,  # Use default font
+                size=text_height_dxf,
+                align=TextEntityAlignment.LEFT
+            )
+            
+            # Create a translation matrix
+            translation = ezdxf.math.Matrix44.translate(text_x, text_y_dxf, 0)
+            # Apply the translation to each path
+            translated_paths = [p.transform(translation) for p in paths]
+        
+            # Render the paths as splines and polylines
+            path.render_splines_and_polylines(
+                msp, 
+                translated_paths, 
+                dxfattribs={"layer": "ANNOTATION", "color": 7}
+            )
+    
+    # Save the DXF
+    dxf_filepath = os.path.join("./outputs", "out.dxf")
+    doc.saveas(dxf_filepath)
+
+    # ---------------------
+    # 9) For the preview images, draw the polygons and place text similarly
+    # ---------------------
+    draw_polygons_inch(final_polygons_inch, output_img, scaling_factor, processed_size[0], color=(0, 0, 255), thickness=2)
+    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():
+        if boundary_polygon is not None:
+            text_height_cv = 0.75
+            text_x_img = int(((inner_min_x + inner_max_x) / 2.0) / scaling_factor)
+            text_y_in = boundary_polygon.bounds[1] + 0.25
+            text_y_img = int(processed_size[0] - (text_y_in / scaling_factor))
+            org = (text_x_img - int(len(annotation_text.strip()) * 6), text_y_img)
+
+            # Method 2: Use two different thicknesses
+            # Draw thicker outline
+            temp_img = np.zeros_like(output_img)
+
+            cv2.putText(
+                temp_img,
+                annotation_text.strip().upper(),
+                org,
+                cv2.FONT_HERSHEY_SIMPLEX,
+                2,
+                (0, 0, 255),  # Red color
+                4,  # Thicker outline
+                cv2.LINE_AA
+            )
+
+            cv2.putText(
+                temp_img,
+                annotation_text.strip().upper(),
+                org,
+                cv2.FONT_HERSHEY_SIMPLEX,
+                2,
+                (0, 0, 0),  # Black to create hole
+                2,  # Thinner inner part
+                cv2.LINE_AA
+            )
+
+            outline_mask = cv2.cvtColor(temp_img, cv2.COLOR_BGR2GRAY)
+            _, outline_mask = cv2.threshold(outline_mask, 1, 255, cv2.THRESH_BINARY)
+
+            output_img[outline_mask > 0] = temp_img[outline_mask > 0]
+                    
+            cv2.putText(
+                new_outlines,
+                annotation_text.strip().upper(),
+                org,
+                cv2.FONT_HERSHEY_SIMPLEX,
+                2,
+                (0, 0, 255),  # Red color
+                4,  # Thicker outline
+                cv2.LINE_AA
+            )
+            
+            cv2.putText(
+                new_outlines,
+                annotation_text.strip().upper(),
+                org,
+                cv2.FONT_HERSHEY_SIMPLEX,
+                2,
+                (255, 255, 255),  # Inner text in white
+                2,  # Thinner inner part
+                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, offset_unit, finger_clearance, add_boundary, boundary_length, boundary_width, annotation_text):
+        try:
+            return predict(img, offset, offset_unit, finger_clearance, add_boundary, boundary_length, boundary_width, annotation_text)
+        except Exception as e:
+            return None, None, None, None, f"Error: {str(e)}"
+    iface = gr.Interface(
+        fn=gradio_predict,
+        inputs=[
+            gr.Image(label="Input Image"),
+            gr.Number(label="Offset value for Mask", value=0.075),
+            gr.Dropdown(label="Offset Unit", choices=["mm", "inches"], value="inches"),
+            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.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, "inches", "No", "No", 300.0, 200.0, "MyTool"],
+            ["./Test21.jpg", 0.075, "inches", "Yes", "Yes", 300.0, 200.0, "Tool2"]
+        ]
+    )
+    iface.launch(share=True)

coin_det.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf6007ec3d4cd303af4cba2e202f68600a904eb23dfc736b4aa29a215201036b
+size 5490003

requirements.txt

@@ -0,0 +1,10 @@
+transformers
+ultralytics==8.3.9
+ezdxf
+gradio
+pydantic==2.10.6
+kornia
+timm
+einops
+shapely
+gevent==22.10.2

scalingtestupdated.py

@@ -0,0 +1,178 @@
+import cv2
+import numpy as np
+import os
+import argparse
+from typing import Union
+from matplotlib import pyplot as plt
+
+
+class ScalingSquareDetector:
+    def __init__(self, feature_detector="ORB", debug=False):
+        """
+        Initialize the detector with the desired feature matching algorithm.
+        :param feature_detector: "ORB" or "SIFT" (default is "ORB").
+        :param debug: If True, saves intermediate images for debugging.
+        """
+        self.feature_detector = feature_detector
+        self.debug = debug
+        self.detector = self._initialize_detector()
+
+    def _initialize_detector(self):
+        """
+        Initialize the chosen feature detector.
+        :return: OpenCV detector object.
+        """
+        if self.feature_detector.upper() == "SIFT":
+            return cv2.SIFT_create()
+        elif self.feature_detector.upper() == "ORB":
+            return cv2.ORB_create()
+        else:
+            raise ValueError("Invalid feature detector. Choose 'ORB' or 'SIFT'.")
+
+    def find_scaling_square(
+        self, target_image, known_size_mm, roi_margin=30
+    ):
+        """
+        Detect the scaling square in the target image based on the reference image.
+        :param reference_image_path: Path to the reference image of the square.
+        :param target_image_path: Path to the target image containing the square.
+        :param known_size_mm: Physical size of the square in millimeters.
+        :param roi_margin: Margin to expand the ROI around the detected square (in pixels).
+        :return: Scaling factor (mm per pixel).
+        """
+        
+        contours, _ = cv2.findContours(
+            target_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
+        )
+
+        if not contours:
+            raise ValueError("No contours found in the cropped ROI.")
+
+        # # Select the largest square-like contour
+        largest_square = None
+        # largest_square_area = 0
+        # for contour in contours:
+        #     x_c, y_c, w_c, h_c = cv2.boundingRect(contour)
+        #     aspect_ratio = w_c / float(h_c)
+        #     if 0.9 <= aspect_ratio <= 1.1:
+        #         peri = cv2.arcLength(contour, True)
+        #         approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
+        #         if len(approx) == 4:
+        #             area = cv2.contourArea(contour)
+        #             if area > largest_square_area:
+        #                 largest_square = contour
+        #                 largest_square_area = area
+
+        # if largest_square is None:
+        #     raise ValueError("No square-like contour found in the ROI.")
+        for contour in contours:
+            largest_square=contour
+        # Draw the largest contour on the original image
+        target_image_color = cv2.cvtColor(target_image, cv2.COLOR_GRAY2BGR)
+        cv2.drawContours(
+            target_image_color, largest_square, -1, (255, 0, 0), 3
+        )
+
+        # if self.debug:
+        cv2.imwrite("largest_contour.jpg", target_image_color)
+
+        # Calculate the bounding rectangle of the largest contour
+        x, y, w, h = cv2.boundingRect(largest_square)
+        square_width_px = w
+        square_height_px = h
+
+        # Calculate the scaling factor
+        avg_square_size_px = (square_width_px + square_height_px) / 2
+        scaling_factor = known_size_mm / avg_square_size_px  # mm per pixel
+
+        return scaling_factor #, square_height_px, square_width_px, roi_binary
+
+    def draw_debug_images(self, output_folder):
+        """
+        Save debug images if enabled.
+        :param output_folder: Directory to save debug images.
+        """
+        if self.debug:
+            if not os.path.exists(output_folder):
+                os.makedirs(output_folder)
+            debug_images = ["largest_contour.jpg"]
+            for img_name in debug_images:
+                if os.path.exists(img_name):
+                    os.rename(img_name, os.path.join(output_folder, img_name))
+
+
+def calculate_scaling_factor(
+    target_image,
+    reference_obj_size_mm=0.955,
+    feature_detector="ORB",
+    debug=False,
+    roi_margin=30,
+):
+    # Initialize detector
+    detector = ScalingSquareDetector(feature_detector=feature_detector, debug=debug)
+
+    # Find scaling square and calculate scaling factor
+    scaling_factor = detector.find_scaling_square(
+        target_image=target_image,
+        known_size_mm=reference_obj_size_mm,
+        roi_margin=roi_margin,
+    )
+
+    # Save debug images
+    if debug:
+        detector.draw_debug_images("debug_outputs")
+
+    return scaling_factor
+
+
+# Example usage:
+if __name__ == "__main__":
+    import os
+    from PIL import Image
+    from ultralytics import YOLO
+    from app import yolo_detect, shrink_bbox
+    from ultralytics.utils.plotting import save_one_box
+
+    for idx, file in enumerate(os.listdir("./sample_images")):
+        img = np.array(Image.open(os.path.join("./sample_images", file)))
+        img = yolo_detect(img, ['box'])
+        model = YOLO("./last.pt")
+        res = model.predict(img, conf=0.6)
+        
+        box_img = save_one_box(res[0].cpu().boxes.xyxy, im=res[0].orig_img, save=False)
+        # img = shrink_bbox(box_img, 1.20)
+        cv2.imwrite(f"./outputs/{idx}_{file}", box_img)
+        
+        print("File: ",f"./outputs/{idx}_{file}")
+        try:
+
+            scaling_factor = calculate_scaling_factor(
+                target_image=box_img,
+                known_square_size_mm=0.955,
+                feature_detector="ORB",
+                debug=False,
+                roi_margin=90,
+            )
+            # cv2.imwrite(f"./outputs/{idx}_binary_{file}", roi_binary)
+            
+            # Square size in mm
+            # square_size_mm = 0.955
+
+            # # Compute the calculated scaling factors and compare
+            # calculated_scaling_factor = square_size_mm / height_px
+            # discrepancy = abs(calculated_scaling_factor - scaling_factor)
+            # import pprint
+            # pprint.pprint({
+            #     "height_px": height_px,
+            #     "width_px": width_px,
+            #     "given_scaling_factor": scaling_factor,
+            #     "calculated_scaling_factor": calculated_scaling_factor,
+            #     "discrepancy": discrepancy,
+            # })
+
+
+            print(f"Scaling Factor (mm per pixel): {scaling_factor:.6f}")
+        except Exception as e:
+            from traceback import print_exc
+            print(print_exc())
+            print(f"Error: {e}")

u2net.py

@@ -0,0 +1,525 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class REBNCONV(nn.Module):
+    def __init__(self,in_ch=3,out_ch=3,dirate=1):
+        super(REBNCONV,self).__init__()
+
+        self.conv_s1 = nn.Conv2d(in_ch,out_ch,3,padding=1*dirate,dilation=1*dirate)
+        self.bn_s1 = nn.BatchNorm2d(out_ch)
+        self.relu_s1 = nn.ReLU(inplace=True)
+
+    def forward(self,x):
+
+        hx = x
+        xout = self.relu_s1(self.bn_s1(self.conv_s1(hx)))
+
+        return xout
+
+## upsample tensor 'src' to have the same spatial size with tensor 'tar'
+def _upsample_like(src,tar):
+
+    src = F.upsample(src,size=tar.shape[2:],mode='bilinear')
+
+    return src
+
+
+### RSU-7 ###
+class RSU7(nn.Module):#UNet07DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU7,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool5 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv7 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv6d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+        hx = self.pool5(hx5)
+
+        hx6 = self.rebnconv6(hx)
+
+        hx7 = self.rebnconv7(hx6)
+
+        hx6d =  self.rebnconv6d(torch.cat((hx7,hx6),1))
+        hx6dup = _upsample_like(hx6d,hx5)
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6dup,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-6 ###
+class RSU6(nn.Module):#UNet06DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU6,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool4 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv6 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv5d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+        hx = self.pool4(hx4)
+
+        hx5 = self.rebnconv5(hx)
+
+        hx6 = self.rebnconv6(hx5)
+
+
+        hx5d =  self.rebnconv5d(torch.cat((hx6,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-5 ###
+class RSU5(nn.Module):#UNet05DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU5,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool3 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv5 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv4d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+        hx = self.pool3(hx3)
+
+        hx4 = self.rebnconv4(hx)
+
+        hx5 = self.rebnconv5(hx4)
+
+        hx4d = self.rebnconv4d(torch.cat((hx5,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4 ###
+class RSU4(nn.Module):#UNet04DRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.pool1 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=1)
+        self.pool2 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=1)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=2)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=1)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx = self.pool1(hx1)
+
+        hx2 = self.rebnconv2(hx)
+        hx = self.pool2(hx2)
+
+        hx3 = self.rebnconv3(hx)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.rebnconv2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.rebnconv1d(torch.cat((hx2dup,hx1),1))
+
+        return hx1d + hxin
+
+### RSU-4F ###
+class RSU4F(nn.Module):#UNet04FRES(nn.Module):
+
+    def __init__(self, in_ch=3, mid_ch=12, out_ch=3):
+        super(RSU4F,self).__init__()
+
+        self.rebnconvin = REBNCONV(in_ch,out_ch,dirate=1)
+
+        self.rebnconv1 = REBNCONV(out_ch,mid_ch,dirate=1)
+        self.rebnconv2 = REBNCONV(mid_ch,mid_ch,dirate=2)
+        self.rebnconv3 = REBNCONV(mid_ch,mid_ch,dirate=4)
+
+        self.rebnconv4 = REBNCONV(mid_ch,mid_ch,dirate=8)
+
+        self.rebnconv3d = REBNCONV(mid_ch*2,mid_ch,dirate=4)
+        self.rebnconv2d = REBNCONV(mid_ch*2,mid_ch,dirate=2)
+        self.rebnconv1d = REBNCONV(mid_ch*2,out_ch,dirate=1)
+
+    def forward(self,x):
+
+        hx = x
+
+        hxin = self.rebnconvin(hx)
+
+        hx1 = self.rebnconv1(hxin)
+        hx2 = self.rebnconv2(hx1)
+        hx3 = self.rebnconv3(hx2)
+
+        hx4 = self.rebnconv4(hx3)
+
+        hx3d = self.rebnconv3d(torch.cat((hx4,hx3),1))
+        hx2d = self.rebnconv2d(torch.cat((hx3d,hx2),1))
+        hx1d = self.rebnconv1d(torch.cat((hx2d,hx1),1))
+
+        return hx1d + hxin
+
+
+##### U^2-Net ####
+class U2NET(nn.Module):
+
+    def __init__(self,in_ch=3,out_ch=1):
+        super(U2NET,self).__init__()
+
+        self.stage1 = RSU7(in_ch,32,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,32,128)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(128,64,256)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(256,128,512)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(512,256,512)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(512,256,512)
+
+        # decoder
+        self.stage5d = RSU4F(1024,256,512)
+        self.stage4d = RSU4(1024,128,256)
+        self.stage3d = RSU5(512,64,128)
+        self.stage2d = RSU6(256,32,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(128,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(256,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(512,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(512,out_ch,3,padding=1)
+
+        self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        #stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #-------------------- decoder --------------------
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,d1)
+
+        d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
+
+        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)
+
+### U^2-Net small ###
+class U2NETP(nn.Module):
+
+    def __init__(self,in_ch=3,out_ch=1):
+        super(U2NETP,self).__init__()
+
+        self.stage1 = RSU7(in_ch,16,64)
+        self.pool12 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage2 = RSU6(64,16,64)
+        self.pool23 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage3 = RSU5(64,16,64)
+        self.pool34 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage4 = RSU4(64,16,64)
+        self.pool45 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage5 = RSU4F(64,16,64)
+        self.pool56 = nn.MaxPool2d(2,stride=2,ceil_mode=True)
+
+        self.stage6 = RSU4F(64,16,64)
+
+        # decoder
+        self.stage5d = RSU4F(128,16,64)
+        self.stage4d = RSU4(128,16,64)
+        self.stage3d = RSU5(128,16,64)
+        self.stage2d = RSU6(128,16,64)
+        self.stage1d = RSU7(128,16,64)
+
+        self.side1 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side2 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side3 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side4 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side5 = nn.Conv2d(64,out_ch,3,padding=1)
+        self.side6 = nn.Conv2d(64,out_ch,3,padding=1)
+
+        self.outconv = nn.Conv2d(6*out_ch,out_ch,1)
+
+    def forward(self,x):
+
+        hx = x
+
+        #stage 1
+        hx1 = self.stage1(hx)
+        hx = self.pool12(hx1)
+
+        #stage 2
+        hx2 = self.stage2(hx)
+        hx = self.pool23(hx2)
+
+        #stage 3
+        hx3 = self.stage3(hx)
+        hx = self.pool34(hx3)
+
+        #stage 4
+        hx4 = self.stage4(hx)
+        hx = self.pool45(hx4)
+
+        #stage 5
+        hx5 = self.stage5(hx)
+        hx = self.pool56(hx5)
+
+        #stage 6
+        hx6 = self.stage6(hx)
+        hx6up = _upsample_like(hx6,hx5)
+
+        #decoder
+        hx5d = self.stage5d(torch.cat((hx6up,hx5),1))
+        hx5dup = _upsample_like(hx5d,hx4)
+
+        hx4d = self.stage4d(torch.cat((hx5dup,hx4),1))
+        hx4dup = _upsample_like(hx4d,hx3)
+
+        hx3d = self.stage3d(torch.cat((hx4dup,hx3),1))
+        hx3dup = _upsample_like(hx3d,hx2)
+
+        hx2d = self.stage2d(torch.cat((hx3dup,hx2),1))
+        hx2dup = _upsample_like(hx2d,hx1)
+
+        hx1d = self.stage1d(torch.cat((hx2dup,hx1),1))
+
+
+        #side output
+        d1 = self.side1(hx1d)
+
+        d2 = self.side2(hx2d)
+        d2 = _upsample_like(d2,d1)
+
+        d3 = self.side3(hx3d)
+        d3 = _upsample_like(d3,d1)
+
+        d4 = self.side4(hx4d)
+        d4 = _upsample_like(d4,d1)
+
+        d5 = self.side5(hx5d)
+        d5 = _upsample_like(d5,d1)
+
+        d6 = self.side6(hx6)
+        d6 = _upsample_like(d6,d1)
+
+        d0 = self.outconv(torch.cat((d1,d2,d3,d4,d5,d6),1))
+
+        return F.sigmoid(d0), F.sigmoid(d1), F.sigmoid(d2), F.sigmoid(d3), F.sigmoid(d4), F.sigmoid(d5), F.sigmoid(d6)

u2netp.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e7567cde013fb64813973ce6e1ecc25a80c05c3ca7adbc5a54f3c3d90991b854
+size 4683258

yolo11n.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0ebbc80d4a7680d14987a577cd21342b65ecfd94632bd9a8da63ae6417644ee1
+size 5613764

yolov8x-worldv2.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:41e771bfbbb8894dd857f3fef7cac3b3578dffd49fd3547101efa6a606a02a0e
+size 146355704