Upload 4 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)
+
+
+

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)