Update app.py

app.py

@@ -41,6 +41,14 @@ class ReferenceBoxNotDetectedError(Exception):
     """Raised when the reference box cannot be detected in the image"""
     pass
 
+class BoundaryExceedsError(Exception):
+    """Raised when the optional boundary dimensions exceed allowed image dimensions."""
+    pass
+
+class BoundaryOverlapError(Exception):
+    """Raised when the optional boundary dimensions are too small and overlap with the inner contours."""
+    pass
+
 # ---------------------
 # Global Model Initialization with caching and print statements
 # ---------------------
@@ -339,14 +347,21 @@ def save_dxf_spline(inflated_contours, scaling_factor, height, finger_clearance=
             print(f"Skipping contour: {e}")
     return doc, final_polygons_inch
 
-def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width, boundary_unit):
+def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width, offset_unit):
     msp = doc.modelspace()
-    if boundary_unit == "mm":
+    # First, if unit is mm, check if values seem too low (accidental inches) and convert them.
+    if offset_unit.lower() == "mm":
+        if boundary_length < 50:
+            boundary_length = boundary_length * 25.4
+        if boundary_width < 50:
+            boundary_width = boundary_width * 25.4
         boundary_length_in = boundary_length / 25.4
         boundary_width_in = boundary_width / 25.4
     else:
         boundary_length_in = boundary_length
         boundary_width_in = boundary_width
+
+    # Compute bounding box from inner contours.
     min_x = float("inf")
     min_y = float("inf")
     max_x = -float("inf")
@@ -360,6 +375,19 @@ def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width
     if min_x == float("inf"):
         print("No tool polygons found, skipping boundary.")
         return None
+
+    # Calculate inner bounding box dimensions.
+    inner_width = max_x - min_x
+    inner_length = max_y - min_y
+
+    # Define a clearance margin (in inches) required between the inner contours and the boundary.
+    clearance_margin = 0.2  # Adjust this value as needed
+
+    # New check: if the provided boundary dimensions are too small relative to the inner contours, raise an error.
+    if boundary_width_in <= inner_width + clearance_margin or boundary_length_in <= inner_length + clearance_margin:
+        raise BoundaryOverlapError("Error: The specified boundary dimensions are too small and overlap with the inner contours. Please provide larger values.")
+
+    # Compute the boundary rectangle centered on the inner contours.
     shape_cx = (min_x + max_x) / 2
     shape_cy = (min_y + max_y) / 2
     half_w = boundary_width_in / 2.0
@@ -369,6 +397,7 @@ def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width
     bottom = shape_cy - half_l
     top = shape_cy + half_l
     rect_coords = [(left, bottom), (right, bottom), (right, top), (left, top), (left, bottom)]
+
     from shapely.geometry import Polygon as ShapelyPolygon
     boundary_polygon = ShapelyPolygon(rect_coords)
     msp.add_lwpolyline(rect_coords, close=True, dxfattribs={"layer": "BOUNDARY"})
@@ -399,12 +428,12 @@ def draw_single_polygon(poly, image_rgb, scaling_factor, image_height, color=(0,
 # ---------------------
 def predict(
     image: Union[str, bytes, np.ndarray],
-    offset_inches: float,
+    offset_value: float,
+    offset_unit: str,         # "mm" or "inches"
     finger_clearance: str,    # "Yes" or "No"
     add_boundary: str,        # "Yes" or "No"
     boundary_length: float,
     boundary_width: float,
-    boundary_unit: str,
     annotation_text: str
 ):
     overall_start = time.time()
@@ -417,7 +446,7 @@ def predict(
                 image = np.array(Image.open(io.BytesIO(base64.b64decode(image))).convert("RGB"))
             except Exception:
                 raise ValueError("Invalid base64 image data")
-    # Apply sharpness enhancement if image is a NumPy array.
+    # Apply sharpness enhancement.
     if isinstance(image, np.ndarray):
         pil_image = Image.fromarray(image)
         enhanced_image = ImageEnhance.Sharpness(pil_image).enhance(1.5)
@@ -462,6 +491,35 @@ def predict(
         print("Using default scaling factor of 1.0 due to calculation error")
     gc.collect()
     print("Scaling factor determined: {}".format(scaling_factor))
+    
+    # ---------------------
+    # Process boundary dimensions if boundary is enabled.
+    # ---------------------
+    if add_boundary.lower() == "yes":
+        image_height_px, image_width_px = shrunked_img.shape[:2]
+        image_height_in = image_height_px * scaling_factor
+        image_width_in = image_width_px * scaling_factor
+        # First, if units are mm, check if boundary dimensions seem too low and convert.
+        if offset_unit.lower() == "mm":
+            if boundary_length < 50:
+                boundary_length = boundary_length * 25.4
+            if boundary_width < 50:
+                boundary_width = boundary_width * 25.4
+            boundary_length_in = boundary_length / 25.4
+            boundary_width_in = boundary_width / 25.4
+        else:
+            boundary_length_in = boundary_length
+            boundary_width_in = boundary_width
+        if boundary_length_in > (image_height_in - 2) or boundary_width_in > (image_width_in - 2):
+            raise BoundaryExceedsError("Error: The specified boundary dimensions exceed the allowed image dimensions (image size minus 2 inches). Please enter smaller values.")
+    
+    # Convert offset value.
+    if offset_unit.lower() == "mm":
+        if offset_value < 1:
+            offset_value = offset_value * 25.4
+        offset_inches = offset_value / 25.4
+    else:
+        offset_inches = offset_value
     t = time.time()
     orig_size = shrunked_img.shape[:2]
     objects_mask = remove_bg(shrunked_img)
@@ -477,17 +535,13 @@ def predict(
     del objects_mask
     gc.collect()
     print("Mask dilation completed in {:.2f} seconds".format(time.time() - t))
-    # Save the dilated mask for debugging if needed.
     Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
-    # --- Extract outlines (only used for DXF generation) ---
     t = time.time()
     outlines, contours = extract_outlines(dilated_mask)
     print("Outline extraction completed in {:.2f} seconds".format(time.time() - t))
-    # Instead of drawing the original contours, we now prepare a clean copy of the shrunk image for drawing new contours.
     output_img = shrunked_img.copy()
     del shrunked_img
     gc.collect()
-    # --- Generate DXF using the extracted contours and apply finger clearance ---
     t = time.time()
     use_finger_clearance = True if finger_clearance.lower() == "yes" else False
     doc, final_polygons_inch = save_dxf_spline(contours, scaling_factor, processed_size[0], finger_clearance=use_finger_clearance)
@@ -496,10 +550,9 @@ def predict(
     print("DXF generation completed in {:.2f} seconds".format(time.time() - t))
     boundary_polygon = None
     if add_boundary.lower() == "yes":
-        boundary_polygon = add_rectangular_boundary(doc, final_polygons_inch, boundary_length, boundary_width, boundary_unit)
+        boundary_polygon = add_rectangular_boundary(doc, final_polygons_inch, boundary_length, boundary_width, offset_unit)
         if boundary_polygon is not None:
             final_polygons_inch.append(boundary_polygon)
-    # --- Annotation Text Placement (Centered horizontally) ---
     min_x = float("inf")
     min_y = float("inf")
     max_x = -float("inf")
@@ -514,31 +567,34 @@ def predict(
             max_x = b[2]
         if b[3] > max_y:
             max_y = b[3]
-    margin = 0.5
+    margin = 0.40
     text_x = (min_x + max_x) / 2
-    text_y = min_y - margin
+    if add_boundary.lower() == "yes":
+        text_y = min_y + margin
+    else:
+        text_y = min_y - margin
     msp = doc.modelspace()
     if annotation_text.strip():
         text_entity = msp.add_text(
             annotation_text.strip(),
             dxfattribs={
-                "height": 0.25,
-                "layer": "ANNOTATION"
+                "height": 0.50,
+                "layer": "ANNOTATION",
+                "style": "Bold"
             }
         )
         text_entity.dxf.insert = (text_x, text_y)
     dxf_filepath = os.path.join("./outputs", "out.dxf")
     doc.saveas(dxf_filepath)
-    # --- Draw only the new contours (final_polygons_inch) on the clean output image ---
     draw_polygons_inch(final_polygons_inch, output_img, scaling_factor, processed_size[0], color=(0,0,255), thickness=2)
-    # Also prepare an "Outlines" image based on a blank canvas for clarity.
     new_outlines = np.ones_like(output_img) * 255
     draw_polygons_inch(final_polygons_inch, new_outlines, scaling_factor, processed_size[0], color=(0,0,255), thickness=2)
     if annotation_text.strip():
         text_px = int(text_x / scaling_factor)
         text_py = int(processed_size[0] - (text_y / scaling_factor))
-        cv2.putText(output_img, annotation_text.strip(), (text_px, text_py), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2, cv2.LINE_AA)
-        cv2.putText(new_outlines, annotation_text.strip(), (text_px, text_py), cv2.FONT_HERSHEY_SIMPLEX, 1, (0,0,255), 2, cv2.LINE_AA)
+        org = (int(text_px) - int(len(annotation_text.strip()) / 2), int(text_py))
+        cv2.putText(output_img, annotation_text.strip(), org, cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0,0,255), 3, cv2.LINE_AA)
+        cv2.putText(new_outlines, annotation_text.strip(), org, cv2.FONT_HERSHEY_SIMPLEX, 1.25, (0,0,255), 3, cv2.LINE_AA)
     outlines_color = cv2.cvtColor(new_outlines, cv2.COLOR_BGR2RGB)
     print("Total prediction time: {:.2f} seconds".format(time.time() - overall_start))
     return (
@@ -554,18 +610,21 @@ def predict(
 # ---------------------
 if __name__ == "__main__":
     os.makedirs("./outputs", exist_ok=True)
-    def gradio_predict(img, offset, finger_clearance, add_boundary, boundary_length, boundary_width, boundary_unit, annotation_text):
-        return predict(img, offset, finger_clearance, add_boundary, boundary_length, boundary_width, boundary_unit, annotation_text)
+    def gradio_predict(img, offset, offset_unit, finger_clearance, add_boundary, boundary_length, boundary_width, annotation_text):
+        try:
+            return predict(img, offset, offset_unit, finger_clearance, add_boundary, boundary_length, boundary_width, annotation_text)
+        except Exception as e:
+            return None, None, None, None, f"Error: {str(e)}"
     iface = gr.Interface(
         fn=gradio_predict,
         inputs=[
             gr.Image(label="Input Image"),
-            gr.Number(label="Offset value for Mask (inches)", value=0.075),
+            gr.Number(label="Offset value for Mask", value=0.075),
+            gr.Dropdown(label="Offset Unit", choices=["mm", "inches"], value="inches"),
             gr.Dropdown(label="Add Finger Clearance?", choices=["Yes", "No"], value="No"),
             gr.Dropdown(label="Add Rectangular Boundary?", choices=["Yes", "No"], value="No"),
             gr.Number(label="Boundary Length", value=300.0, precision=2),
             gr.Number(label="Boundary Width", value=200.0, precision=2),
-            gr.Dropdown(label="Boundary Unit", choices=["mm", "inches"], value="mm"),
             gr.Textbox(label="Annotation (max 20 chars)", max_length=20, placeholder="Type up to 20 characters")
         ],
         outputs=[
@@ -576,8 +635,8 @@ if __name__ == "__main__":
             gr.Textbox(label="Scaling Factor (inches/pixel)")
         ],
         examples=[
-            ["./Test20.jpg", 0.075, "No", "No", 300.0, 200.0, "mm", "MyTool"],
-            ["./Test21.jpg", 0.075, "Yes", "Yes", 300.0, 200.0, "mm", "Tool2"]
+            ["./Test20.jpg", 0.075, "inches", "No", "No", 300.0, 200.0, "MyTool"],
+            ["./Test21.jpg", 0.075, "inches", "Yes", "Yes", 300.0, 200.0, "Tool2"]
         ]
     )
     iface.launch(share=True)