Upload 6 files

api_server.py

@@ -0,0 +1,525 @@
+# from fastapi import FastAPI, HTTPException, UploadFile, File, Form
+# from pydantic import BaseModel
+# import numpy as np
+# from PIL import Image
+# import io, uuid, os, shutil, timeit
+# from datetime import datetime
+# from fastapi.staticfiles import StaticFiles
+# from fastapi.middleware.cors import CORSMiddleware
+
+# # import your three wrappers
+# from app import predict_simple, predict_middle, predict_full
+
+# app = FastAPI()
+
+# # allow CORS if needed
+# app.add_middleware(
+#     CORSMiddleware,
+#     allow_origins=["*"],
+#     allow_methods=["*"],
+#     allow_headers=["*"],
+# )
+
+# BASE_URL = "https://snapanddtraceapp-988917236820.us-central1.run.app"
+# OUTPUT_DIR = os.path.abspath("./outputs")
+# os.makedirs(OUTPUT_DIR, exist_ok=True)
+# app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")
+
+# UPDATES_DIR = os.path.abspath("./updates")
+# os.makedirs(UPDATES_DIR, exist_ok=True)
+# app.mount("/updates", StaticFiles(directory=UPDATES_DIR), name="updates")
+
+
+# def save_and_build_urls(
+#     session_id: str,
+#     output_image: np.ndarray,
+#     outlines: np.ndarray,
+#     dxf_path: str,
+#     mask: np.ndarray
+# ):
+#     """Helper to save all four artifacts and return public URLs."""
+#     request_dir = os.path.join(OUTPUT_DIR, session_id)
+#     os.makedirs(request_dir, exist_ok=True)
+
+#     # filenames
+#     out_fn = "overlay.jpg"
+#     outlines_fn = "outlines.jpg"
+#     mask_fn = "mask.jpg"
+#     current_date = datetime.now().strftime("%d-%m-%Y")
+#     dxf_fn = f"out_{current_date}_{session_id}.dxf"
+
+#     # full paths
+#     out_path = os.path.join(request_dir, out_fn)
+#     outlines_path = os.path.join(request_dir, outlines_fn)
+#     mask_path = os.path.join(request_dir, mask_fn)
+#     new_dxf_path = os.path.join(request_dir, dxf_fn)
+
+#     # save images
+#     Image.fromarray(output_image).save(out_path)
+#     Image.fromarray(outlines).save(outlines_path)
+#     Image.fromarray(mask).save(mask_path)
+
+#     # copy dx file
+#     if os.path.exists(dxf_path):
+#         shutil.copy(dxf_path, new_dxf_path)
+#     else:
+#         # fallback if your DXF generator returns bytes or string
+#         with open(new_dxf_path, "wb") as f:
+#             if isinstance(dxf_path, (bytes, bytearray)):
+#                 f.write(dxf_path)
+#             else:
+#                 f.write(str(dxf_path).encode("utf-8"))
+
+#     # build URLs
+#     return {
+#         "output_image_url": f"{BASE_URL}/outputs/{session_id}/{out_fn}",
+#         "outlines_url":      f"{BASE_URL}/outputs/{session_id}/{outlines_fn}",
+#         "mask_url":          f"{BASE_URL}/outputs/{session_id}/{mask_fn}",
+#         "dxf_url":           f"{BASE_URL}/outputs/{session_id}/{dxf_fn}",
+#     }
+
+
+# @app.post("/predict1")
+# async def predict1_api(
+#     file: UploadFile = File(...)
+# ):
+#     """
+#     Simple predict: only image → overlay, outlines, mask, DXF
+#     """
+#     session_id = str(uuid.uuid4())
+#     try:
+#         img_bytes = await file.read()
+#         image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
+#     except Exception:
+#         raise HTTPException(400, "Invalid image upload")
+
+#     try:
+#         start = timeit.default_timer()
+#         out_img, outlines, dxf_path, mask = predict_simple(image)
+#         elapsed = timeit.default_timer() - start
+#         print(f"[{session_id}] predict1 in {elapsed:.2f}s")
+
+#         return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)
+
+#     except Exception as e:
+#         raise HTTPException(500, f"predict1 failed: {e}")
+#     except ReferenceBoxNotDetectedError:
+#         raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
+#     except FingerCutOverlapError:
+#         raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")
+
+
+# @app.post("/predict2")
+# async def predict2_api(
+#     file: UploadFile = File(...),
+#     enable_fillet: str = Form(..., regex="^(On|Off)$"),
+#     fillet_value_mm: float = Form(...)
+# ):
+#     """
+#     Middle predict: image + fillet toggle + fillet value → overlay, outlines, mask, DXF
+#     """
+#     session_id = str(uuid.uuid4())
+#     try:
+#         img_bytes = await file.read()
+#         image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
+#     except Exception:
+#         raise HTTPException(400, "Invalid image upload")
+
+#     try:
+#         start = timeit.default_timer()
+#         out_img, outlines, dxf_path, mask = predict_middle(
+#             image, enable_fillet, fillet_value_mm
+#         )
+#         elapsed = timeit.default_timer() - start
+#         print(f"[{session_id}] predict2 in {elapsed:.2f}s")
+
+#         return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)
+
+#     except Exception as e:
+#         raise HTTPException(500, f"predict2 failed: {e}")
+#     except ReferenceBoxNotDetectedError:
+#         raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
+#     except FingerCutOverlapError:
+#         raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")
+
+# @app.post("/predict3")
+# async def predict3_api(
+#     file: UploadFile = File(...),
+#     enable_fillet: str = Form(..., regex="^(On|Off)$"),
+#     fillet_value_mm: float = Form(...),
+#     enable_finger_cut: str = Form(..., regex="^(On|Off)$")
+# ):
+#     """
+#     Full predict: image + fillet toggle/value + finger-cut toggle → overlay, outlines, mask, DXF
+#     """
+#     session_id = str(uuid.uuid4())
+#     try:
+#         img_bytes = await file.read()
+#         image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
+#     except Exception:
+#         raise HTTPException(400, "Invalid image upload")
+
+#     try:
+#         start = timeit.default_timer()
+#         out_img, outlines, dxf_path, mask = predict_full(
+#             image, enable_fillet, fillet_value_mm, enable_finger_cut
+#         )
+#         elapsed = timeit.default_timer() - start
+#         print(f"[{session_id}] predict3 in {elapsed:.2f}s")
+
+#         return save_and_build_urls(session_id, out_img, outlines, dxf_path, mask)
+
+#     except Exception as e:
+#         raise HTTPException(500, f"predict3 failed: {e}")
+#     except ReferenceBoxNotDetectedError:
+#         raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
+#     except FingerCutOverlapError:
+#         raise HTTPException(status_code=400, detail="There was an overlap with fingercuts!s Please try again to generate dxf.")
+
+# @app.post("/update")
+# async def update_files(
+#     output_image: UploadFile = File(...),
+#     outlines_image: UploadFile = File(...),
+#     mask_image: UploadFile = File(...),
+#     dxf_file: UploadFile = File(...)
+# ):
+#     session_id = str(uuid.uuid4())
+#     update_dir = os.path.join(UPDATES_DIR, session_id)
+#     os.makedirs(update_dir, exist_ok=True)
+
+#     try:
+#         upload_map = {
+#             "output_image":  output_image,
+#             "outlines_image": outlines_image,
+#             "mask_image":     mask_image,
+#             "dxf_file":       dxf_file,
+#         }
+#         urls = {}
+#         for key, up in upload_map.items():
+#             fn = up.filename
+#             path = os.path.join(update_dir, fn)
+#             with open(path, "wb") as f:
+#                 shutil.copyfileobj(up.file, f)
+#             urls[key] = f"{BASE_URL}/updates/{session_id}/{fn}"
+
+#         return {"session_id": session_id, "uploaded": urls}
+
+#     except Exception as e:
+#         raise HTTPException(500, f"Update failed: {e}")
+
+
+# if __name__ == "__main__":
+#     import uvicorn
+#     port = int(os.environ.get("PORT", 8082))
+#     print(f"Starting FastAPI server on 0.0.0.0:{port}...")
+#     uvicorn.run(app, host="0.0.0.0", port=port)
+
+
+
+
+
+
+
+
+
+
+from fastapi import FastAPI, HTTPException, UploadFile, File, Form
+from pydantic import BaseModel
+import numpy as np
+from PIL import Image
+import io, uuid, os, shutil, timeit
+from datetime import datetime
+from fastapi.staticfiles import StaticFiles
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import FileResponse
+
+# import your three wrappers
+from app import predict_simple, predict_middle, predict_full
+
+from app import (
+    predict_simple, predict_middle, predict_full,
+    ReferenceBoxNotDetectedError,
+    FingerCutOverlapError
+)
+
+
+app = FastAPI()
+
+# allow CORS if needed
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+BASE_URL = "https://snapanddtraceapp-988917236820.us-central1.run.app"
+
+OUTPUT_DIR = os.path.abspath("./outputs")
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+
+UPDATES_DIR = os.path.abspath("./updates")
+os.makedirs(UPDATES_DIR, exist_ok=True)
+
+# Mount static directories with normal StaticFiles
+app.mount("/outputs", StaticFiles(directory=OUTPUT_DIR), name="outputs")
+app.mount("/updates", StaticFiles(directory=UPDATES_DIR), name="updates")
+
+
+def save_and_build_urls(
+    session_id: str,
+    output_image: np.ndarray,
+    outlines: np.ndarray,
+    dxf_path: str,
+    mask: np.ndarray,
+    endpoint_type: str,
+    fillet_value: float = None,
+    finger_cut: str = None
+):
+    """Helper to save all four artifacts and return public URLs."""
+    request_dir = os.path.join(OUTPUT_DIR, session_id)
+    os.makedirs(request_dir, exist_ok=True)
+
+    # filenames
+    out_fn = "overlay.jpg"
+    outlines_fn = "outlines.jpg"
+    mask_fn = "mask.jpg"
+    
+    # Get current date
+    current_date = datetime.utcnow().strftime("%d-%m-%Y")
+
+
+    # Format fillet value with underscore instead of dot
+    fillet_str = f"{fillet_value:.2f}".replace(".", "_") if fillet_value is not None else None
+    
+    # Determine DXF filename based on endpoint type
+    if endpoint_type == "predict1":
+        dxf_fn = f"DXF_{current_date}.dxf"
+    elif endpoint_type == "predict2":
+        dxf_fn = f"DXF_{current_date}.dxf"
+    elif endpoint_type == "predict3":
+        dxf_fn = f"DXF_{current_date}.dxf"
+
+    # full paths
+    out_path = os.path.join(request_dir, out_fn)
+    outlines_path = os.path.join(request_dir, outlines_fn)
+    mask_path = os.path.join(request_dir, mask_fn)
+    new_dxf_path = os.path.join(request_dir, dxf_fn)
+
+    # save images
+    Image.fromarray(output_image).save(out_path)
+    Image.fromarray(outlines).save(outlines_path)
+    Image.fromarray(mask).save(mask_path)
+
+    # copy dxf file
+    if os.path.exists(dxf_path):
+        shutil.copy(dxf_path, new_dxf_path)
+    else:
+        # fallback if your DXF generator returns bytes or string
+        with open(new_dxf_path, "wb") as f:
+            if isinstance(dxf_path, (bytes, bytearray)):
+                f.write(dxf_path)
+            else:
+                f.write(str(dxf_path).encode("utf-8"))
+
+    # build URLs with /download prefix for DXF
+    return {
+        "output_image_url": f"{BASE_URL}/outputs/{session_id}/{out_fn}",
+        "outlines_url":     f"{BASE_URL}/outputs/{session_id}/{outlines_fn}",
+        "mask_url":         f"{BASE_URL}/outputs/{session_id}/{mask_fn}",
+        "dxf_url":          f"{BASE_URL}/download/{session_id}/{dxf_fn}",  # Changed to use download endpoint
+    }
+
+# Add new endpoint for downloading DXF files
+@app.get("/download/{session_id}/{filename}")
+async def download_file(session_id: str, filename: str):
+    file_path = os.path.join(OUTPUT_DIR, session_id, filename)
+    if not os.path.exists(file_path):
+        raise HTTPException(status_code=404, detail="File not found")
+    
+    return FileResponse(
+        path=file_path,
+        filename=filename,
+        media_type="application/x-dxf",
+        headers={"Content-Disposition": f"attachment; filename={filename}"}
+    )
+
+
+@app.post("/predict1")
+async def predict1_api(
+    file: UploadFile = File(...)
+):
+    """
+    Simple predict: only image → overlay, outlines, mask, DXF
+    DXF naming format: DXF_DD-MM-YYYY.dxf
+    """
+    session_id = str(uuid.uuid4())
+    try:
+        img_bytes = await file.read()
+        image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
+    except Exception:
+        raise HTTPException(400, "Invalid image upload")
+
+    try:
+        start = timeit.default_timer()
+        out_img, outlines, dxf_path, mask = predict_simple(image)
+        elapsed = timeit.default_timer() - start
+        print(f"[{session_id}] predict1 in {elapsed:.2f}s")
+
+        return save_and_build_urls(
+            session_id=session_id,
+            output_image=out_img,
+            outlines=outlines,
+            dxf_path=dxf_path,
+            mask=mask,
+            endpoint_type="predict1"
+        )
+
+    except ReferenceBoxNotDetectedError:
+        raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
+    except FingerCutOverlapError:
+        raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")
+
+@app.post("/predict2")
+async def predict2_api(
+    file: UploadFile = File(...),
+    enable_fillet: str = Form(..., regex="^(On|Off)$"),
+    fillet_value_mm: float = Form(...)
+):
+    """
+    Middle predict: image + fillet toggle + fillet value → overlay, outlines, mask, DXF
+    DXF naming format: DXF_DD-MM-YYYY_fillet-value_mm.dxf
+    """
+    session_id = str(uuid.uuid4())
+    try:
+        img_bytes = await file.read()
+        image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
+    except Exception:
+        raise HTTPException(400, "Invalid image upload")
+
+    try:
+        start = timeit.default_timer()
+        out_img, outlines, dxf_path, mask = predict_middle(
+            image, enable_fillet, fillet_value_mm
+        )
+        elapsed = timeit.default_timer() - start
+        print(f"[{session_id}] predict2 in {elapsed:.2f}s")
+
+        return save_and_build_urls(
+            session_id=session_id,
+            output_image=out_img,
+            outlines=outlines,
+            dxf_path=dxf_path,
+            mask=mask,
+            endpoint_type="predict2",
+            fillet_value=fillet_value_mm
+        )
+
+    except ReferenceBoxNotDetectedError:
+        raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
+    except FingerCutOverlapError:
+        raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")
+
+
+@app.post("/predict3")
+async def predict3_api(
+    file: UploadFile = File(...),
+    enable_fillet: str = Form(..., regex="^(On|Off)$"),
+    fillet_value_mm: float = Form(...),
+    enable_finger_cut: str = Form(..., regex="^(On|Off)$")
+):
+    """
+    Full predict: image + fillet toggle/value + finger-cut toggle → overlay, outlines, mask, DXF
+    DXF naming format: DXF_DD-MM-YYYY_fillet-value_mm_fingercut-On|Off.dxf
+    """
+    session_id = str(uuid.uuid4())
+    try:
+        img_bytes = await file.read()
+        image = np.array(Image.open(io.BytesIO(img_bytes)).convert("RGB"))
+    except Exception:
+        raise HTTPException(400, "Invalid image upload")
+
+    try:
+        start = timeit.default_timer()
+        out_img, outlines, dxf_path, mask = predict_full(
+            image, enable_fillet, fillet_value_mm, enable_finger_cut
+        )
+        elapsed = timeit.default_timer() - start
+        print(f"[{session_id}] predict3 in {elapsed:.2f}s")
+
+        return save_and_build_urls(
+            session_id=session_id,
+            output_image=out_img,
+            outlines=outlines,
+            dxf_path=dxf_path,
+            mask=mask,
+            endpoint_type="predict3",
+            fillet_value=fillet_value_mm,
+            finger_cut=enable_finger_cut
+        )
+
+    except ReferenceBoxNotDetectedError:
+        raise HTTPException(status_code=400, detail="Error detecting reference battery! Please try again with a clearer image.")
+    except FingerCutOverlapError:
+        raise HTTPException(status_code=400, detail="There was an overlap with fingercuts! Please try again to generate dxf.")
+    except HTTPException as e:
+        raise e
+    except Exception as e:
+        raise HTTPException(status_code=500, detail="Error detecting reference battery! Please try again with a clearer image.")
+
+
+@app.post("/update")
+async def update_files(
+    output_image: UploadFile = File(...),
+    outlines_image: UploadFile = File(...),
+    mask_image: UploadFile = File(...),
+    dxf_file: UploadFile = File(...)
+):
+    session_id = str(uuid.uuid4())
+    update_dir = os.path.join(UPDATES_DIR, session_id)
+    os.makedirs(update_dir, exist_ok=True)
+
+    try:
+        upload_map = {
+            "output_image":  output_image,
+            "outlines_image": outlines_image,
+            "mask_image":     mask_image,
+            "dxf_file":       dxf_file,
+        }
+        urls = {}
+        for key, up in upload_map.items():
+            fn = up.filename
+            path = os.path.join(update_dir, fn)
+            with open(path, "wb") as f:
+                shutil.copyfileobj(up.file, f)
+            urls[key] = f"{BASE_URL}/updates/{session_id}/{fn}"
+
+        return {"session_id": session_id, "uploaded": urls}
+
+    except Exception as e:
+        raise HTTPException(500, f"Update failed: {e}")
+
+
+from fastapi import Response
+
+@app.get("/health")
+def health():
+    return Response(content="OK", status_code=200)
+
+
+if __name__ == "__main__":
+    import uvicorn
+    port = int(os.environ.get("PORT", 8080))
+    print(f"Starting FastAPI server on 0.0.0.0:{port}...")
+    uvicorn.run(app, host="0.0.0.0", port=port)
+
+
+

app.py.txt

@@ -0,0 +1,1007 @@
+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 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
+
+# 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)
+
+# 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")
+
+# 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")
+
+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):
+            paper_detector_global = YOLO(paper_model_path)
+        else:
+            # Fallback to generic object detection for paper-like rectangles
+            logger.warning("Using fallback paper detection")
+            paper_detector_global = None
+        logger.info("Paper detector loaded successfully")
+    return paper_detector_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) -> Tuple[np.ndarray, float]:
+    """
+    Detect paper in the image using contour detection as fallback
+    Returns the paper contour and estimated scaling factor
+    """
+    # Convert to grayscale
+    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if len(image.shape) == 3 else image
+    
+    # Apply Gaussian blur
+    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
+    
+    # Edge detection
+    edges = cv2.Canny(blurred, 50, 150)
+    
+    # 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 = []
+    min_area = (image.shape[0] * image.shape[1]) * 0.1  # At least 10% of image
+    
+    for contour in contours:
+        area = cv2.contourArea(contour)
+        if area > min_area:
+            # Approximate contour to polygon
+            epsilon = 0.02 * cv2.arcLength(contour, True)
+            approx = cv2.approxPolyDP(contour, epsilon, True)
+            
+            # Check if it's roughly rectangular (4 corners)
+            if len(approx) >= 4:
+                # Calculate bounding rectangle
+                rect = cv2.boundingRect(approx)
+                aspect_ratio = rect[2] / rect[3]  # width / height
+                
+                # Check if aspect ratio matches common paper ratios
+                # A4: 1.414, A3: 1.414, US Letter: 1.294
+                if 0.7 < aspect_ratio < 1.8:  # Allow some tolerance
+                    paper_contours.append((contour, area, aspect_ratio))
+    
+    if not paper_contours:
+        raise PaperNotDetectedError("Could not detect paper in the image")
+    
+    # Select the largest paper-like contour
+    paper_contours.sort(key=lambda x: x[1], reverse=True)
+    best_contour = paper_contours[0][0]
+    
+    return best_contour, 0.0  # Return 0.0 as placeholder scaling factor
+
+def detect_paper_bounds(image: np.ndarray, paper_size: str) -> 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.5)
+            if not results or len(results) == 0 or len(results[0].boxes) == 0:
+                logger.warning("Model detection failed, using fallback contour detection")
+                return detect_paper_contour(image)
+            
+            # Get the largest detected paper
+            boxes = results[0].cpu().boxes.xyxy
+            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:
+                raise PaperNotDetectedError("No paper detected by model")
+                
+            # 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]]
+            ])
+            
+        else:
+            # Use fallback contour detection
+            paper_contour, _ = detect_paper_contour(image)
+        
+        # Calculate scaling factor based on paper size
+        scaling_factor = calculate_paper_scaling_factor(paper_contour, paper_size)
+        
+        return paper_contour, scaling_factor
+        
+    except Exception as e:
+        logger.error(f"Error in paper detection: {e}")
+        raise PaperNotDetectedError(f"Failed to detect paper: {str(e)}")
+
+def calculate_paper_scaling_factor(paper_contour: np.ndarray, paper_size: str) -> float:
+    """
+    Calculate scaling factor based on detected paper dimensions
+    """
+    # Get paper dimensions
+    paper_dims = PAPER_SIZES[paper_size]
+    expected_width_mm = paper_dims["width"]
+    expected_height_mm = paper_dims["height"]
+    
+    # Calculate bounding rectangle of paper contour
+    rect = cv2.boundingRect(paper_contour)
+    detected_width_px = rect[2]
+    detected_height_px = rect[3]
+    
+    # Calculate scaling factors for both dimensions
+    scale_x = expected_width_mm / detected_width_px
+    scale_y = expected_height_mm / detected_height_px
+    
+    # Use average of both scales
+    scaling_factor = (scale_x + scale_y) / 2
+    
+    logger.info(f"Paper detection: {detected_width_px}x{detected_height_px} px -> {expected_width_mm}x{expected_height_mm} mm")
+    logger.info(f"Calculated scaling factor: {scaling_factor:.4f} mm/px")
+    
+    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)
+    min_area = 1000  # Minimum area threshold
+    significant_contours = [c for c in contours if cv2.contourArea(c) > min_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 exclude_paper_area(mask: np.ndarray, paper_contour: np.ndarray, expansion_factor: float = 1.1) -> np.ndarray:
+    """
+    Remove paper area from the mask to focus only on objects
+    """
+    # Create paper mask with slight expansion to ensure complete removal
+    paper_mask = np.zeros(mask.shape[:2], dtype=np.uint8)
+    
+    # Expand paper contour slightly
+    epsilon = expansion_factor * cv2.arcLength(paper_contour, True)
+    expanded_contour = cv2.approxPolyDP(paper_contour, epsilon, True)
+    
+    cv2.fillPoly(paper_mask, [expanded_contour], 255)
+    
+    # Invert paper mask and apply to object mask
+    paper_mask_inv = cv2.bitwise_not(paper_mask)
+    result_mask = cv2.bitwise_and(mask, paper_mask_inv)
+    
+    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"""
+    doc = ezdxf.new(units=ezdxf.units.MM)
+    doc.header["$INSUNITS"] = ezdxf.units.MM
+    msp = doc.modelspace()
+    final_polygons_inch = []
+    finger_centers = []
+    original_polygons = []
+
+    # Scale correction factor
+    scale_correction = 1.079
+
+    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
+
+            tool_polygon = build_tool_polygon(points_inch)
+            original_polygons.append(tool_polygon)
+
+            if finger_clearance:
+                try:
+                    tool_polygon, center = place_finger_cut_adjusted(
+                        tool_polygon, points_inch, finger_centers, final_polygons_inch
+                    )
+                except FingerCutOverlapError:
+                    tool_polygon = original_polygons[-1]
+
+            exterior_coords = polygon_to_exterior_coords(tool_polygon)
+            if len(exterior_coords) < 3:
+                continue
+
+            # Apply scale correction
+            corrected_coords = [(x * scale_correction, y * scale_correction) for x, y in exterior_coords]
+
+            msp.add_spline(corrected_coords, degree=3, dxfattribs={"layer": "TOOLS"})
+            final_polygons_inch.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_inch, 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 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, edge_radius, finger_clearance=False):
+    """Main prediction function using paper as reference"""
+    
+    if offset_unit == "inches":
+        offset *= 25.4
+
+    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
+        paper_contour, scaling_factor = detect_paper_bounds(image, paper_size)
+        logger.info(f"Paper detected with scaling factor: {scaling_factor:.4f} mm/px")
+        
+    except PaperNotDetectedError as e:
+        return (
+            None, None, None, None,
+            f"Error: {str(e)}"
+        )
+    except Exception as e:
+        raise gr.Error(f"Error processing image: {str(e)}")
+
+    try:
+        # Remove background from main objects
+        orig_size = image.shape[:2]
+        objects_mask = remove_bg(image)
+        processed_size = objects_mask.shape[:2]
+        
+        # Resize mask to match original image
+        objects_mask = cv2.resize(objects_mask, (image.shape[1], image.shape[0]))
+        
+        # Remove paper area from mask to focus only on objects
+        objects_mask = exclude_paper_area(objects_mask, paper_contour)
+        
+        # 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
+    if edge_radius > 0:
+        rounded_mask = round_edges(objects_mask, edge_radius, scaling_factor)
+    else:
+        rounded_mask = objects_mask.copy()
+    
+    # Apply dilation for offset
+    if offset > 0:
+        offset_pixels = (float(offset) / scaling_factor) * 2 + 1 if scaling_factor else 1
+        kernel = np.ones((int(offset_pixels), int(offset_pixels)), np.uint8)
+        dilated_mask = cv2.dilate(rounded_mask, kernel)
+    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
+        dxf, finger_polygons, original_polygons = save_dxf_spline(
+            contours,
+            scaling_factor,
+            processed_size[0],
+            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(processed_size[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)
+    
+    # Draw paper bounds on annotated image
+    cv2.drawContours(shrunked_img_contours, [paper_contour], -1, (255, 0, 0), thickness=3)
+    
+    # Add paper size text
+    paper_text = f"Paper: {paper_size}"
+    cv2.putText(shrunked_img_contours, paper_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2)
+    
+    cleanup_models()
+
+    return (
+        shrunked_img_contours,
+        outlines,
+        dxf,
+        dilated_mask_orig,
+        f"Scale: {scaling_factor:.4f} mm/px | Paper: {paper_size}"
+    )
+
+def predict_full_paper(image, paper_size, enable_fillet, fillet_value_mm, enable_finger_cut, selected_outputs):
+    """
+    Full prediction function with paper reference and flexible outputs
+    Returns DXF + conditionally selected additional outputs
+    """
+    radius = fillet_value_mm if enable_fillet == "On" else 0
+    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=0,  # No offset for now, can be added as parameter later
+        offset_unit="mm",
+        edge_radius=radius,
+        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:
+        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.
+        
+        **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
+        """)
+        
+        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("### Edge Rounding")
+                    enable_fillet = gr.Radio(
+                        choices=["On", "Off"],
+                        value="Off",
+                        label="Enable Edge Rounding",
+                        interactive=True
+                    )
+                    
+                    fillet_value_mm = gr.Slider(
+                        minimum=0,
+                        maximum=20,
+                        step=1,
+                        value=5,
+                        label="Edge Radius (mm)",
+                        visible=False,
+                        interactive=True
+                    )
+
+                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)
+        
+        # Dynamic visibility updates
+        def toggle_fillet(choice):
+            return gr.update(visible=(choice == "On"))
+
+        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)
+            ]
+        
+        # Event handlers
+        enable_fillet.change(
+            fn=toggle_fillet,
+            inputs=enable_fillet,
+            outputs=fillet_value_mm
+        )
+        
+        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, 
+                enable_fillet, 
+                fillet_value_mm, 
+                enable_finger_cut, 
+                output_options
+            ],
+            outputs=[dxf_file, output_image, outlines_image, mask_image, scale_info]
+        )
+        
+        # Example gallery
+        with gr.Row():
+            gr.Markdown("""
+            ### Tips for Best Results:
+            - Ensure good lighting and clear paper edges
+            - Place object completely on the paper
+            - Avoid shadows that might interfere with detection
+            - Use high contrast between object and paper
+            """)
+
+    demo.launch(share=True)

requirements.txt

@@ -0,0 +1,13 @@
+fastapi
+transformers==4.48.3
+ultralytics==8.3.9
+pydantic==2.10.6
+ezdxf==1.3.5
+gradio==5.15.0
+kornia==0.8.0
+timm==1.0.14
+einops==0.8.1
+torchvision==0.20.1
+torch==2.5.1
+torchaudio==2.5.1
+shapely

scalingtestupdated.py

@@ -0,0 +1,184 @@
+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
+        print(f"No of contours: {len(contours)}")
+        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
+
+        for contour in contours:
+            largest_square = contour
+
+        # if largest_square is None:
+        #     raise ValueError("No square-like contour found in the ROI.")
+
+        # 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
+        print(f"Reference object size:  {known_size_mm} mm")
+        print(f"width:  {square_width_px} px")
+        print(f"height:  {square_height_px} px")
+
+        # Calculate the scaling factor
+        avg_square_size_px = (square_width_px + square_height_px) / 2
+        print(f"avg square size: {avg_square_size_px} px")
+        scaling_factor = known_size_mm / avg_square_size_px  # mm per pixel
+        print(f"scaling factor: {scaling_factor} 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,
+    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("./best.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=20,
+                feature_detector="ORB",
+                debug=False,
+                roi_margin=90,
+            )
+            # cv2.imwrite(f"./outputs/{idx}_binary_{file}", roi_binary)
+            
+            # Square size in mm
+            # square_size_mm = 12.7
+
+            # # 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}")

u2netp.pth

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

u2netp.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)