app.py

import os
from pathlib import Path
from typing import List, Union, Tuple
from PIL import Image
import ezdxf.units
import numpy as np
import torch
from torchvision import transforms
from ultralytics import YOLOWorld, YOLO
from ultralytics.engine.results import Results
from ultralytics.utils.plotting import save_one_box
from transformers import AutoModelForImageSegmentation
import cv2
import ezdxf
import gradio as gr
import gc
# from scalingtestupdated import calculate_scaling_factor
from scalingtestupdated import calculate_scaling_factor_with_units, calculate_paper_scaling_factor, convert_units, calculate_paper_scaling_factor_corrected
from scipy.interpolate import splprep, splev
from scipy.ndimage import gaussian_filter1d
import json
import time
import signal
from shapely.ops import unary_union
from shapely.geometry import MultiPolygon, GeometryCollection, Polygon, Point
from u2netp import U2NETP
import logging
import shutil
import sys

# Add this at the very beginning of your main Python file, before any other imports
os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '0'
os.environ['OPENCV_IO_ENABLE_JASPER'] = '0'
os.environ['QT_QPA_PLATFORM'] = 'offscreen'
os.environ['MPLBACKEND'] = 'Agg'

# For headless environments
import matplotlib
matplotlib.use('Agg')

# Initialize logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Create cache directory for models
CACHE_DIR = os.path.join(os.path.dirname(__file__), ".cache")
os.makedirs(CACHE_DIR, exist_ok=True)

# Paper size configurations (in mm)
PAPER_SIZES = {
    "A4": {"width": 210, "height": 297},
    "A3": {"width": 297, "height": 420},
    "US Letter": {"width": 215.9, "height": 279.4}
}
    
# Custom Exception Classes
class TimeoutReachedError(Exception):
    pass

class BoundaryOverlapError(Exception):
    pass

class TextOverlapError(Exception):
    pass

class PaperNotDetectedError(Exception):
    """Raised when the paper cannot be detected in the image"""
    pass

class MultipleObjectsError(Exception):
    """Raised when multiple objects are detected on the paper"""
    def __init__(self, message="Multiple objects detected. Please place only a single object on the paper."):
        super().__init__(message)

class NoObjectDetectedError(Exception):
    """Raised when no object is detected on the paper"""
    def __init__(self, message="No object detected on the paper. Please ensure an object is placed on the paper."):
        super().__init__(message)

class FingerCutOverlapError(Exception):
    """Raised when finger cuts overlap with existing geometry"""
    def __init__(self, message="There was an overlap with fingercuts... Please try again to generate dxf."):
        super().__init__(message)

class ReferenceBoxNotDetectedError(Exception):
    """Raised when reference box/paper cannot be detected"""
    def __init__(self, message="Reference box not detected"):
        super().__init__(message)

# Global model variables for lazy loading
paper_detector_global = None
u2net_global = None
birefnet = None

# Model paths
paper_model_path = os.path.join(CACHE_DIR, "paper_detector.pt")  # You'll need to train/provide this
u2net_model_path = os.path.join(CACHE_DIR, "u2netp.pth")

# Global variable for YOLOWorld
yolo_v8_global = None
yolo_v8_model_path = os.path.join(CACHE_DIR, "yolov8n.pt")  # Adjust path as needed

    
# Device configuration
device = "cpu"
torch.set_float32_matmul_precision(["high", "highest"][0])

def ensure_model_files():
    """Ensure model files are available in cache directory"""
    if not os.path.exists(paper_model_path):
        if os.path.exists("paper_detector.pt"):
            shutil.copy("paper_detector.pt", paper_model_path)
        else:
            logger.warning("paper_detector.pt model file not found - using fallback detection")
    
    if not os.path.exists(u2net_model_path):
        if os.path.exists("u2netp.pth"):
            shutil.copy("u2netp.pth", u2net_model_path)
        else:
            raise FileNotFoundError("u2netp.pth model file not found")
    logger.info("YOLOv8 will auto-download if not present")

ensure_model_files()

# Lazy loading functions
def get_paper_detector():
    """Lazy load paper detector model"""
    global paper_detector_global
    if paper_detector_global is None:
        logger.info("Loading paper detector model...")
        if os.path.exists(paper_model_path):
            try:
                paper_detector_global = YOLO(paper_model_path)
                logger.info("Paper detector loaded successfully")
            except Exception as e:
                logger.error(f"Failed to load paper detector: {e}")
                paper_detector_global = None
        else:
            # Fallback to generic object detection for paper-like rectangles
            logger.warning("Paper model file not found, using fallback detection")
            paper_detector_global = None
    return paper_detector_global
def get_yolo_v8():
    """Lazy load YOLOv8 model"""
    global yolo_v8_global
    if yolo_v8_global is None:
        logger.info("Loading YOLOv8 model...")
        try:
            yolo_v8_global = YOLO(yolo_v8_model_path)  # Auto-downloads if needed
            logger.info("YOLOv8 model loaded successfully")
        except Exception as e:
            logger.error(f"Failed to load YOLOv8: {e}")
            yolo_v8_global = None
    return yolo_v8_global
def get_u2net():
    """Lazy load U2NETP model"""
    global u2net_global
    if u2net_global is None:
        logger.info("Loading U2NETP model...")
        u2net_global = U2NETP(3, 1)
        u2net_global.load_state_dict(torch.load(u2net_model_path, map_location="cpu"))
        u2net_global.to(device)
        u2net_global.eval()
        logger.info("U2NETP model loaded successfully")
    return u2net_global

def load_birefnet_model():
    """Load BiRefNet model from HuggingFace"""
    return AutoModelForImageSegmentation.from_pretrained(
        'ZhengPeng7/BiRefNet', 
        trust_remote_code=True
    )

def get_birefnet():
    """Lazy load BiRefNet model"""
    global birefnet
    if birefnet is None:
        logger.info("Loading BiRefNet model...")
        birefnet = load_birefnet_model()
        birefnet.to(device)
        birefnet.eval()
        logger.info("BiRefNet model loaded successfully")
    return birefnet

def detect_paper_contour(image: np.ndarray, output_unit: str = "mm") -> Tuple[np.ndarray, float]:
    """
    Detect paper in the image using contour detection as fallback
    Returns the paper contour and estimated scaling factor
    """
    logger.info("Using contour-based paper detection")
    
    # Convert to grayscale
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
    
    # Apply bilateral filter to reduce noise while preserving edges
    filtered = cv2.bilateralFilter(gray, 9, 75, 75)
    
    # Apply adaptive threshold
    thresh = cv2.adaptiveThreshold(filtered, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
                                   cv2.THRESH_BINARY, 11, 2)
    
    # Edge detection with multiple thresholds
    edges1 = cv2.Canny(filtered, 50, 150)
    edges2 = cv2.Canny(filtered, 30, 100)
    edges = cv2.bitwise_or(edges1, edges2)
    
    # Morphological operations to connect broken edges
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3))
    edges = cv2.morphologyEx(edges, cv2.MORPH_CLOSE, kernel)
    
    # Find contours
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    # Filter contours by area and aspect ratio to find paper-like rectangles
    paper_contours = []
    image_area = image.shape[0] * image.shape[1]
    min_area = image_area * 0.20  # At least 15% of image
    max_area = image_area * 0.85  # At most 95% of image
    
    for contour in contours:
        area = cv2.contourArea(contour)
        if min_area < area < max_area:
            # Approximate contour to polygon
            epsilon = 0.015 * cv2.arcLength(contour, True)
            approx = cv2.approxPolyDP(contour, epsilon, True)
            
            # Check if it's roughly rectangular (4 corners) or close to it
            if len(approx) >= 4:
                # Calculate bounding rectangle
                rect = cv2.boundingRect(approx)
                w, h = rect[2], rect[3]
                aspect_ratio = w / h if h > 0 else 0
                
                # Check if aspect ratio matches common paper ratios
                # A4: 1.414, A3: 1.414, US Letter: 1.294
                if 1.3 < aspect_ratio < 1.5:  # More lenient tolerance
                    # Check if contour area is close to bounding rect area (rectangularity)
                    rect_area = w * h
                    if rect_area > 0:
                        extent = area / rect_area
                        if extent > 0.85:  # At least 85% rectangular
                            paper_contours.append((contour, area, aspect_ratio, extent))
    
    if not paper_contours:
        logger.error("No paper-like contours found")
        raise ReferenceBoxNotDetectedError("Could not detect paper in the image using contour detection")
    
    # Select the best paper contour based on area and rectangularity
    paper_contours.sort(key=lambda x: (x[1] * x[3]), reverse=True)  # Sort by area * extent
    best_contour = paper_contours[0][0]
    
    logger.info(f"Paper detected using contours: area={paper_contours[0][1]}, aspect_ratio={paper_contours[0][2]:.2f}")
    
    # Return 0.0 as placeholder - will be calculated later based on paper size
    return best_contour, 0.0
    
def detect_paper_bounds(image: np.ndarray, paper_size: str, output_unit: str = "mm") -> Tuple[np.ndarray, float]:
    """
    Detect paper bounds in the image and calculate scaling factor
    """
    try:
        paper_detector = get_paper_detector()
        
        if paper_detector is not None:
            # Use trained model if available
            results = paper_detector.predict(image, conf=0.8, verbose=False)
            
            if not results or len(results) == 0:
                logger.warning("No results from paper detector")
                return detect_paper_contour(image, output_unit)
                
            # Check if boxes exist and are not empty
            if not hasattr(results[0], 'boxes') or results[0].boxes is None or len(results[0].boxes) == 0:
                logger.warning("No boxes detected by model, using fallback contour detection")
                return detect_paper_contour(image, output_unit)
            
            # Get the largest detected paper
            boxes = results[0].boxes.xyxy.cpu().numpy()
            if len(boxes) == 0:
                logger.warning("Empty boxes detected, using fallback")
                return detect_paper_contour(image, output_unit)
                
            largest_box = None
            max_area = 0
            
            for box in boxes:
                x_min, y_min, x_max, y_max = box
                area = (x_max - x_min) * (y_max - y_min)
                if area > max_area:
                    max_area = area
                    largest_box = box
            
            if largest_box is None:
                logger.warning("No valid paper box found, using fallback")
                return detect_paper_contour(image, output_unit)
                
            # Convert box to contour-like format
            x_min, y_min, x_max, y_max = map(int, largest_box)
            paper_contour = np.array([
                [[x_min, y_min]],
                [[x_max, y_min]],
                [[x_max, y_max]],
                [[x_min, y_max]]
            ])
            
            logger.info(f"Paper detected by model: {x_min},{y_min} to {x_max},{y_max}")
            
        else:
            # Use fallback contour detection
            logger.info("Using fallback contour detection for paper")
            paper_contour, _ = detect_paper_contour(image, output_unit)

        # After getting paper_contour, expand it
        rect = cv2.boundingRect(paper_contour)
        expansion = int(min(rect[2], rect[3]) * 0.1)  # Expand by 10%
        
        x, y, w, h = rect
        expanded_contour = np.array([
            [[max(0, x - expansion), max(0, y - expansion)]],
            [[min(image.shape[1], x + w + expansion), max(0, y - expansion)]],
            [[min(image.shape[1], x + w + expansion), min(image.shape[0], y + h + expansion)]],
            [[max(0, x - expansion), min(image.shape[0], y + h + expansion)]]
        ])
        
        paper_contour = expanded_contour
        
        # Calculate scaling factor based on paper size with proper units
        # scaling_factor = calculate_paper_scaling_factor(paper_contour, paper_size, output_unit)
        scaling_factor, unit_string = calculate_paper_scaling_factor_corrected(
            paper_contour, 
            paper_size, 
            output_unit="mm",
            correction_factor=1.235,  # Adjust this value
            method="average"  # Try different methods
)        
        return paper_contour, scaling_factor
        
    except Exception as e:
        logger.error(f"Error in paper detection: {e}")
        raise ReferenceBoxNotDetectedError(f"Failed to detect paper: {str(e)}")

def calculate_paper_scaling_factor(paper_contour: np.ndarray, paper_size: str, output_unit: str = "mm") -> float:
    """
    Calculate scaling factor based on detected paper dimensions with proper unit handling.
    """
    from scalingtestupdated import calculate_paper_scaling_factor as calc_paper_scale
    scaling_factor, unit_string = calc_paper_scale(paper_contour, paper_size, output_unit)
    return scaling_factor


def validate_single_object(mask: np.ndarray, paper_contour: np.ndarray) -> None:
    """
    Validate that only a single object is present on the paper
    """
    # Create a mask for the paper area
    paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8)
    cv2.fillPoly(paper_mask, [paper_contour], 255)
    
    # Apply paper mask to object mask
    masked_objects = cv2.bitwise_and(mask, paper_mask)
    
    # Find contours of objects within paper bounds
    contours, _ = cv2.findContours(masked_objects, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
   
    # Filter out very small contours (noise) and paper-sized contours
    image_area = mask.shape[0] * mask.shape[1]
    min_area = 100  # Minimum area threshold
    max_area = image_area * 0.5  # Maximum 50% of image area (to exclude paper detection)
    significant_contours = [c for c in contours if min_area < cv2.contourArea(c) < max_area]
        
    if len(significant_contours) == 0:
        raise NoObjectDetectedError()
    elif len(significant_contours) > 1:
        raise MultipleObjectsError()
    
    logger.info(f"Single object validated: {len(significant_contours)} significant contour(s) found")

def remove_bg_u2netp(image: np.ndarray) -> np.ndarray:
    """Remove background using U2NETP model"""
    try:
        u2net_model = get_u2net()
        
        image_pil = Image.fromarray(image)
        transform_u2netp = transforms.Compose([
            transforms.Resize((320, 320)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])
        
        input_tensor = transform_u2netp(image_pil).unsqueeze(0).to(device)
        
        with torch.no_grad():
            outputs = u2net_model(input_tensor)
        
        pred = outputs[0]
        pred = (pred - pred.min()) / (pred.max() - pred.min() + 1e-8)
        pred_np = pred.squeeze().cpu().numpy()
        pred_np = cv2.resize(pred_np, (image_pil.width, image_pil.height))
        pred_np = (pred_np * 255).astype(np.uint8)
        
        return pred_np
    except Exception as e:
        logger.error(f"Error in U2NETP background removal: {e}")
        raise


def remove_bg(image: np.ndarray) -> np.ndarray:
    """Remove background using BiRefNet model for main objects"""
    try:
        birefnet_model = get_birefnet()
        
        transform_image = transforms.Compose([
            transforms.Resize((1024, 1024)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
        ])
        
        image_pil = Image.fromarray(image)
        input_images = transform_image(image_pil).unsqueeze(0).to(device)

        with torch.no_grad():
            preds = birefnet_model(input_images)[-1].sigmoid().cpu()
        pred = preds[0].squeeze()

        pred_pil = transforms.ToPILImage()(pred)
        
        scale_ratio = 1024 / max(image_pil.size)
        scaled_size = (int(image_pil.size[0] * scale_ratio), int(image_pil.size[1] * scale_ratio))
        
        return np.array(pred_pil.resize(scaled_size))
    except Exception as e:
        logger.error(f"Error in BiRefNet background removal: {e}")
        raise

def mask_paper_area_in_image(image: np.ndarray, paper_contour: np.ndarray) -> np.ndarray:
    """Less aggressive masking to preserve corner objects"""
    masked_image = image.copy()
    
    # Much less aggressive shrinking - only 2% instead of 8%
    rect = cv2.boundingRect(paper_contour)
    shrink_pixels = max(5, int(min(rect[2], rect[3]) * 0.02))  # Changed from 0.08 to 0.02
    
    x, y, w, h = rect
    # Create mask but keep more area
    outer_mask = np.ones(image.shape[:2], dtype=np.uint8) * 255
    
    inner_contour = np.array([
        [[x + shrink_pixels, y + shrink_pixels]],
        [[x + w - shrink_pixels, y + shrink_pixels]],
        [[x + w - shrink_pixels, y + h - shrink_pixels]],
        [[x + shrink_pixels, y + h - shrink_pixels]]
    ])
    
    cv2.fillPoly(outer_mask, [inner_contour], 0)
    masked_image[outer_mask == 255] = [128, 128, 128]  # Gray instead of black
    
    return masked_image
    
def exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.2) -> np.ndarray:
    """Less aggressive paper area exclusion"""
    # Create paper mask
    paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8)
    cv2.fillPoly(paper_mask, [paper_contour], 255)
    
    # Instead of eroding, slightly expand the paper mask
    rect = cv2.boundingRect(paper_contour)
    expansion = max(10, int(min(rect[2], rect[3]) * 0.02))  # 2% expansion
    
    kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (expansion, expansion))
    expanded_paper_mask = cv2.dilate(paper_mask, kernel, iterations=1)
    
    # Keep objects within expanded paper area
    result_mask = cv2.bitwise_and(mask, expanded_paper_mask)
    
    return result_mask

def resample_contour(contour, edge_radius_px: int = 0):
    """Resample contour with radius-aware smoothing and periodic handling."""
    logger.info(f"Starting resample_contour with contour of shape {contour.shape}")

    num_points = 1500
    sigma = max(2, int(edge_radius_px) // 4)

    if len(contour) < 4:
        error_msg = f"Contour must have at least 4 points, but has {len(contour)} points."
        logger.error(error_msg)
        raise ValueError(error_msg)

    try:
        contour = contour[:, 0, :]
        logger.debug(f"Reshaped contour to shape {contour.shape}")

        if not np.array_equal(contour[0], contour[-1]):
            contour = np.vstack([contour, contour[0]])

        tck, u = splprep(contour.T, u=None, s=0, per=True)
        
        u_new = np.linspace(u.min(), u.max(), num_points)
        x_new, y_new = splev(u_new, tck, der=0)
        
        if sigma > 0:
            x_new = gaussian_filter1d(x_new, sigma=sigma, mode='wrap')
            y_new = gaussian_filter1d(y_new, sigma=sigma, mode='wrap')
            
            x_new[-1] = x_new[0]
            y_new[-1] = y_new[0]

        result = np.array([x_new, y_new]).T
        logger.info(f"Completed resample_contour with result shape {result.shape}")
        return result

    except Exception as e:
        logger.error(f"Error in resample_contour: {e}")
        raise


def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=False):
    """Save contours as DXF splines with optional finger cuts - scaling_factor should be in mm/px"""
    doc = ezdxf.new(units=ezdxf.units.MM)
    doc.header["$INSUNITS"] = ezdxf.units.MM
    msp = doc.modelspace()
    final_polygons_mm = []  # Use mm instead of inch for clarity
    finger_centers = []
    original_polygons = []

    for contour in inflated_contours:
        try:
            resampled_contour = resample_contour(contour)  

            # Convert pixel coordinates to mm using the scaling factor
            points_mm = [(x * scaling_factor, (height - y) * scaling_factor) 
                        for x, y in resampled_contour]

            if len(points_mm) < 3:
                continue

            tool_polygon = build_tool_polygon(points_mm)
            original_polygons.append(tool_polygon)

            if finger_clearance:
                try:
                    tool_polygon, center = place_finger_cut_adjusted(
                        tool_polygon, points_mm, finger_centers, final_polygons_mm
                    )
                except FingerCutOverlapError:
                    tool_polygon = original_polygons[-1]

            exterior_coords = polygon_to_exterior_coords(tool_polygon)
            if len(exterior_coords) < 3:
                continue

            # Coordinates are already in mm, so add directly to DXF
            msp.add_spline(exterior_coords, degree=3, dxfattribs={"layer": "TOOLS"})
            final_polygons_mm.append(tool_polygon)

        except ValueError as e:
            logger.warning(f"Skipping contour: {e}")

    dxf_filepath = os.path.join("./outputs", "out.dxf")
    doc.saveas(dxf_filepath)
    return dxf_filepath, final_polygons_mm, original_polygons

def build_tool_polygon(points_inch):
    """Build a polygon from inch-converted points"""
    return Polygon(points_inch)

def polygon_to_exterior_coords(poly):
    """Extract exterior coordinates from polygon"""
    logger.info(f"Starting polygon_to_exterior_coords with input geometry type: {poly.geom_type}")

    try:
        if poly.geom_type == "GeometryCollection" or poly.geom_type == "MultiPolygon":
            logger.debug(f"Performing unary_union on {poly.geom_type}")
            unified = unary_union(poly)
            if unified.is_empty:
                logger.warning("unary_union produced an empty geometry; returning empty list")
                return []
            
            if unified.geom_type == "GeometryCollection" or unified.geom_type == "MultiPolygon":
                largest = None
                max_area = 0.0
                for g in getattr(unified, "geoms", []):
                    if hasattr(g, "area") and g.area > max_area and hasattr(g, "exterior"):
                        max_area = g.area
                        largest = g
                if largest is None:
                    logger.warning("No valid Polygon found in unified geometry; returning empty list")
                    return []
                poly = largest
            else:
                poly = unified

        if not hasattr(poly, "exterior") or poly.exterior is None:
            logger.warning("Input geometry has no exterior ring; returning empty list")
            return []

        raw_coords = list(poly.exterior.coords)
        total = len(raw_coords)
        logger.info(f"Extracted {total} raw exterior coordinates")

        if total == 0:
            return []

        # Subsample coordinates to at most 100 points
        max_pts = 100
        if total > max_pts:
            step = total // max_pts
            sampled = [raw_coords[i] for i in range(0, total, step)]
            if sampled[-1] != raw_coords[-1]:
                sampled.append(raw_coords[-1])
            logger.info(f"Downsampled perimeter from {total} to {len(sampled)} points")
            return sampled
        else:
            return raw_coords

    except Exception as e:
        logger.error(f"Error in polygon_to_exterior_coords: {e}")
        return []

def place_finger_cut_adjusted(
    tool_polygon: Polygon,
    points_inch: list,
    existing_centers: list,
    all_polygons: list,
    circle_diameter: float = 25.4,
    min_gap: float = 0.5,
    max_attempts: int = 100
) -> Tuple[Polygon, tuple]:
    """Place finger cuts with collision avoidance"""
    logger.info(f"Starting place_finger_cut_adjusted with {len(points_inch)} input points")

    def fallback_solution():
        logger.warning("Using fallback approach for finger cut placement")
        fallback_center = points_inch[len(points_inch) // 2]
        r = circle_diameter / 2.0
        fallback_circle = Point(fallback_center).buffer(r, resolution=32)
        try:
            union_poly = tool_polygon.union(fallback_circle)
        except Exception as e:
            logger.warning(f"Fallback union failed ({e}); trying buffer-union fallback")
            union_poly = tool_polygon.buffer(0).union(fallback_circle.buffer(0))

        existing_centers.append(fallback_center)
        logger.info(f"Fallback finger cut placed at {fallback_center}")
        return union_poly, fallback_center

    r = circle_diameter / 2.0
    needed_center_dist = circle_diameter + min_gap

    raw_perimeter = polygon_to_exterior_coords(tool_polygon)
    if not raw_perimeter:
        logger.warning("No valid exterior coords found; using fallback immediately")
        return fallback_solution()

    if len(raw_perimeter) > 100:
        step = len(raw_perimeter) // 100
        perimeter_coords = raw_perimeter[::step]
        logger.info(f"Subsampled perimeter from {len(raw_perimeter)} to {len(perimeter_coords)} points")
    else:
        perimeter_coords = raw_perimeter[:]

    indices = list(range(len(perimeter_coords)))
    np.random.shuffle(indices)
    logger.debug(f"Shuffled perimeter indices for candidate order")

    start_time = time.time()
    timeout_secs = 5.0

    attempts = 0
    try:
        while attempts < max_attempts:
            if time.time() - start_time > timeout_secs - 0.1:
                logger.warning(f"Approaching timeout after {attempts} attempts")
                return fallback_solution()

            for idx in indices:
                if time.time() - start_time > timeout_secs - 0.05:
                    logger.warning("Timeout during candidate-point loop")
                    return fallback_solution()

                cx, cy = perimeter_coords[idx]
                for dx, dy in [(0, 0), (-min_gap/2, 0), (min_gap/2, 0), (0, -min_gap/2), (0, min_gap/2)]:
                    candidate_center = (cx + dx, cy + dy)

                    # Check distance to existing finger centers
                    too_close_finger = any(
                        np.hypot(candidate_center[0] - ex, candidate_center[1] - ey) 
                        < needed_center_dist 
                        for (ex, ey) in existing_centers
                    )
                    if too_close_finger:
                        continue

                    # Build candidate circle
                    candidate_circle = Point(candidate_center).buffer(r, resolution=32)

                    # Must overlap ≥30% with this polygon
                    try:
                        inter_area = tool_polygon.intersection(candidate_circle).area
                    except Exception:
                        continue

                    if inter_area < 0.3 * candidate_circle.area:
                        continue

                    # Must not intersect other polygons
                    invalid = False
                    for other_poly in all_polygons:
                        if other_poly.equals(tool_polygon):
                            continue
                        if other_poly.buffer(min_gap).intersects(candidate_circle) or \
                           other_poly.buffer(min_gap).touches(candidate_circle):
                            invalid = True
                            break
                    if invalid:
                        continue

                    # Union and return
                    try:
                        union_poly = tool_polygon.union(candidate_circle)
                        if union_poly.geom_type == "MultiPolygon" and len(union_poly.geoms) > 1:
                            continue
                        if union_poly.equals(tool_polygon):
                            continue
                    except Exception:
                        continue

                    existing_centers.append(candidate_center)
                    logger.info(f"Finger cut placed successfully at {candidate_center} after {attempts} attempts")
                    return union_poly, candidate_center

            attempts += 1
            if attempts >= (max_attempts // 2) and (time.time() - start_time) > timeout_secs * 0.8:
                logger.warning(f"Approaching timeout (attempt {attempts})")
                return fallback_solution()

        logger.warning(f"No valid spot after {max_attempts} attempts, using fallback")
        return fallback_solution()

    except Exception as e:
        logger.error(f"Error in place_finger_cut_adjusted: {e}")
        return fallback_solution()

# def extract_outlines(binary_image: np.ndarray) -> Tuple[np.ndarray, list]:
#     """Extract outlines from binary image"""
#     contours, _ = cv2.findContours(
#         binary_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
#     )
#     outline_image = np.full_like(binary_image, 255)
#     return outline_image, contours

def extract_outlines(binary_image: np.ndarray) -> Tuple[np.ndarray, list]:
    """Extract outlines from binary image"""
    # Check if contours are being cut at image boundaries
    h, w = binary_image.shape
    
    # Add small border to prevent boundary cutting
    bordered_image = cv2.copyMakeBorder(binary_image, 5, 5, 5, 5, cv2.BORDER_CONSTANT, value=0)
    
    contours, _ = cv2.findContours(
        bordered_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE
    )
    
    # Adjust contour coordinates back to original image space
    adjusted_contours = []
    for contour in contours:
        adjusted_contour = contour - [5, 5]  # Subtract border offset
        adjusted_contours.append(adjusted_contour)
    
    outline_image = np.full_like(binary_image, 255)
    return outline_image, adjusted_contours

def round_edges(mask: np.ndarray, radius_mm: float, scaling_factor: float) -> np.ndarray:
    """Round mask edges using contour smoothing"""
    if radius_mm <= 0 or scaling_factor <= 0:
        return mask
    
    radius_px = max(1, int(radius_mm / scaling_factor))
    
    if np.count_nonzero(mask) < 500:
        return cv2.dilate(cv2.erode(mask, np.ones((3,3))), np.ones((3,3)))
    
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours = [c for c in contours if cv2.contourArea(c) > 100]
    smoothed_contours = []
    
    for contour in contours:
        try:
            resampled = resample_contour(contour, radius_px)
            resampled = resampled.astype(np.int32).reshape((-1, 1, 2))
            smoothed_contours.append(resampled)
        except Exception as e:
            logger.warning(f"Error smoothing contour: {e}")
            smoothed_contours.append(contour)
    
    rounded = np.zeros_like(mask)
    cv2.drawContours(rounded, smoothed_contours, -1, 255, thickness=cv2.FILLED)
    
    return rounded

def cleanup_memory():
    """Clean up memory after processing"""
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()
    logger.info("Memory cleanup completed")

def cleanup_models():
    """Unload models to free memory"""
    global paper_detector_global, u2net_global, birefnet
    if paper_detector_global is not None:
        del paper_detector_global
        paper_detector_global = None
    if u2net_global is not None:
        del u2net_global
        u2net_global = None
    if birefnet is not None:
        del birefnet
        birefnet = None
    cleanup_memory()

def make_square(img: np.ndarray):
    """Make the image square by padding"""
    height, width = img.shape[:2]
    max_dim = max(height, width)
    
    pad_height = (max_dim - height) // 2
    pad_width = (max_dim - width) // 2
    
    pad_height_extra = max_dim - height - 2 * pad_height
    pad_width_extra = max_dim - width - 2 * pad_width
    
    if len(img.shape) == 3:
        padded = np.pad(
            img,
            (
                (pad_height, pad_height + pad_height_extra),
                (pad_width, pad_width + pad_width_extra),
                (0, 0),
            ),
            mode="edge",
        )
    else:
        padded = np.pad(
            img,
            (
                (pad_height, pad_height + pad_height_extra),
                (pad_width, pad_width + pad_width_extra),
            ),
            mode="edge",
        )
    
    return padded

def predict_with_paper(image, paper_size, offset, offset_unit, finger_clearance=False):
    """Main prediction function using paper as reference"""
    
    logger.info(f"Starting prediction with image shape: {image.shape}")
    logger.info(f"Paper size: {paper_size}, Offset: {offset} {offset_unit}")
    
    # Convert offset to mm for internal calculations (DXF generation expects mm)
    if offset_unit == "inches":
        offset_mm = convert_units(offset, "inches", "mm")
    else:
        offset_mm = offset
        
    edge_radius = None
    if edge_radius is None or edge_radius == 0:
        edge_radius = 0.0001

    if offset < 0:
        raise gr.Error("Offset Value Can't be negative")

    try:
        # Detect paper bounds and calculate scaling factor (always in mm for DXF)
        logger.info("Starting paper detection...")
        paper_contour, scaling_factor = detect_paper_bounds(image, paper_size, output_unit="mm")
        logger.info(f"Paper detected successfully with scaling factor: {scaling_factor:.6f} mm/px")
        
    except ReferenceBoxNotDetectedError as e:
        logger.error(f"Paper detection failed: {e}")
        return (
            None, None, None, None,
            f"Error: {str(e)}"
        )
    except Exception as e:
        logger.error(f"Unexpected error in paper detection: {e}")
        raise gr.Error(f"Error processing image: {str(e)}")

    try:
            # Get paper bounds with expansion
            rect = cv2.boundingRect(paper_contour)
            expansion = max(20, int(min(rect[2], rect[3]) * 0.05))  # 5% expansion
            
            x, y, w, h = rect
            x_min = max(0, x - expansion)
            y_min = max(0, y - expansion)
            x_max = min(image.shape[1], x + w + expansion)
            y_max = min(image.shape[0], y + h + expansion)
            
            # Process the expanded paper area
            cropped_image = image[y_min:y_max, x_min:x_max]
            crop_offset = (x_min, y_min)
            
            # Remove background
            objects_mask = remove_bg(cropped_image)
            
            # Resize mask back to cropped image size
            target_height = y_max - y_min
            target_width = x_max - x_min
            objects_mask_resized = cv2.resize(objects_mask, (target_width, target_height))
            
            # Place back in full image space
            full_mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8)
            full_mask[y_min:y_max, x_min:x_max] = objects_mask_resized
    
            # Light filtering only - don't exclude paper area aggressively
            # Just remove small noise
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
            objects_mask = cv2.morphologyEx(full_mask, cv2.MORPH_OPEN, kernel)
            
            # Debug: Save intermediate masks
            cv2.imwrite("./debug/objects_mask_after_processing.jpg", objects_mask)
            
            # Check if we actually have object pixels
            object_pixels = np.count_nonzero(objects_mask)
            if object_pixels < 300:  # Minimum threshold
                raise NoObjectDetectedError("No significant object detected")
            
            # Validate single object
            validate_single_object(objects_mask, paper_contour)
        
    except (MultipleObjectsError, NoObjectDetectedError) as e:
        return (
            None, None, None, None,
            f"Error: {str(e)}"
        )
    except Exception as e:
        raise gr.Error(f"Error in object detection: {str(e)}")

    # Apply edge rounding if specified
    rounded_mask = objects_mask.copy()
    
    # Apply dilation for offset - more precise calculation using mm values
    if offset_mm > 0:
        offset_pixels = max(1, int(round(float(offset_mm) / scaling_factor)))
        if offset_pixels > 0:
            kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (offset_pixels*2+1, offset_pixels*2+1))
            dilated_mask = cv2.dilate(rounded_mask, kernel, iterations=1)
        else:
            dilated_mask = rounded_mask.copy()
    else:
        dilated_mask = rounded_mask.copy()

    # Save original dilated mask for output
    Image.fromarray(dilated_mask).save("./outputs/scaled_mask_original.jpg")
    dilated_mask_orig = dilated_mask.copy()

    # Extract contours
    outlines, contours = extract_outlines(dilated_mask)

    try:
        # Generate DXF - scaling_factor should be in mm/px for proper DXF units
        dxf, finger_polygons, original_polygons = save_dxf_spline(
            contours,
            scaling_factor,  # This should be mm/px
            image.shape[0],  # Use original image height
            finger_clearance=(finger_clearance == "On")
        )
    except FingerCutOverlapError as e:
        raise gr.Error(str(e))

    # Create annotated image
    shrunked_img_contours = image.copy()

    if finger_clearance == "On":
        outlines = np.full_like(dilated_mask, 255)
        for poly in finger_polygons:
            try:
                coords = np.array([
                    (int(x / scaling_factor), int(image.shape[0] - y / scaling_factor))
                    for x, y in poly.exterior.coords
                ], np.int32).reshape((-1, 1, 2))

                cv2.drawContours(shrunked_img_contours, [coords], -1, (0, 255, 0), thickness=2)
                cv2.drawContours(outlines, [coords], -1, 0, thickness=2)
            except Exception as e:
                logger.warning(f"Failed to draw finger cut: {e}")
                continue
    else:
        outlines = np.full_like(dilated_mask, 255)
        cv2.drawContours(shrunked_img_contours, contours, -1, (0, 255, 0), thickness=2)
        cv2.drawContours(outlines, contours, -1, 0, thickness=2)
    
    cleanup_models()

    # Format scaling info with proper unit display
    if offset_unit == "inches":
        offset_display = f"{offset} inches ({offset_mm:.3f} mm)"
    else:
        offset_display = f"{offset} mm"
    
    scale_info = f"Scale: {scaling_factor:.6f} mm/px | Paper: {paper_size} | Offset: {offset_display}"

    return (
        shrunked_img_contours,
        outlines,
        dxf,
        dilated_mask_orig,
        scale_info
    )

def predict_full_paper(image, paper_size, offset_value_mm = 0.02,offset_unit='mm', enable_finger_cut='Off', selected_outputs=None):
    finger_flag = "On" if enable_finger_cut == "On" else "Off"
    
    # Always get all outputs from predict_with_paper
    ann, outlines, dxf_path, mask, scale_info = predict_with_paper(
        image,
        paper_size,
        offset=offset_value_mm,
        offset_unit= offset_unit,
        finger_clearance=finger_flag,
    )
    
    # Return based on selected outputs
    return (
        dxf_path,  # Always return DXF
        ann if "Annotated Image" in selected_outputs else None,
        outlines if "Outlines" in selected_outputs else None,
        mask if "Mask" in selected_outputs else None,
        scale_info  # Always return scaling info
    )

# Gradio Interface
if __name__ == "__main__":
    os.makedirs("./outputs", exist_ok=True)

    with gr.Blocks(title="Paper-Based DXF Generator", theme=gr.themes.Soft()) as demo:
         # Example gallery
        gr.Markdown("""
            # Paper-Based DXF Generator
            
            Upload an image with a single object placed on paper (A4, A3, or US Letter).
            The paper serves as a size reference for accurate DXF generation.
            """)
        with gr.Row():
            gr.Markdown("""
            **Instructions:**
            1. Place a single object on paper
            2. Select the correct paper size
            3. Configure options as needed
            4. Click Submit to generate DXF
            """)
            gr.Markdown("""
            ### Tips for Best Results:
            - Ensure good lighting and clear paper edges
            - Place object completely at the center of the paper
            - Avoid shadows that might interfere with detection
            - Use high contrast between object and paper
            """)
        
        
        with gr.Row():
            with gr.Column():
                input_image = gr.Image(
                    label="Input Image (Object on Paper)", 
                    type="numpy",
                    height=400
                )
                
                paper_size = gr.Radio(
                    choices=["A4", "A3", "US Letter"],
                    value="A4",
                    label="Paper Size",
                    info="Select the paper size used in your image"
                )
                
                # with gr.Group():
                #     gr.Markdown("### Contour Offset")
                #     with gr.Row():
                #         offset_value_mm = gr.Number(
                #             value=0.02,
                #             label="Offset",
                #             info="Expand contours outward by this amount",
                #             precision=3,
                #             minimum=0,
                #             maximum=50
                #         )
                #         offset_unit = gr.Dropdown(
                #             choices=["mm", "inches"],
                #             value="mm",
                #             label="Unit"
                #         )

                # with gr.Group():
                #     gr.Markdown("### Finger Cuts")
                #     enable_finger_cut = gr.Radio(
                #         choices=["On", "Off"],
                #         value="Off",
                #         label="Enable Finger Cuts",
                #         info="Add circular cuts for easier handling"
                #     )
                
                output_options = gr.CheckboxGroup(
                    choices=["Annotated Image", "Outlines", "Mask"],
                    value=[],
                    label="Additional Outputs",
                    info="DXF is always included"
                )
                
                submit_btn = gr.Button("Generate DXF", variant="primary", size="lg")
            
            with gr.Column():
                with gr.Group():
                    gr.Markdown("### Generated Files")
                    dxf_file = gr.File(label="DXF File", file_types=[".dxf"])
                    scale_info = gr.Textbox(label="Scaling Information", interactive=False)
                
                with gr.Group():
                    gr.Markdown("### Preview Images")
                    output_image = gr.Image(label="Annotated Image", visible=False)
                    outlines_image = gr.Image(label="Outlines", visible=False) 
                    mask_image = gr.Image(label="Mask", visible=False)
        
        def update_outputs_visibility(selected):
            return [
                gr.update(visible="Annotated Image" in selected),
                gr.update(visible="Outlines" in selected),
                gr.update(visible="Mask" in selected)
            ]
        
        
        output_options.change(
            fn=update_outputs_visibility,
            inputs=output_options,
            outputs=[output_image, outlines_image, mask_image]
        )
        
        submit_btn.click(
            fn=predict_full_paper,
            inputs=[
                input_image, 
                paper_size, 
                gr.Number(value=0.02, visible=False),  # Create hidden components
                gr.Textbox(value='mm', visible=False),
                gr.Textbox(value='Off', visible=False),
                output_options
            ],
            outputs=[dxf_file, output_image, outlines_image, mask_image, scale_info]
        )
        
       

    demo.launch(share=True)