Update app.py

app.py

@@ -218,7 +218,7 @@ def shrink_bbox(image: np.ndarray, shrink_factor: float):
     x1 = max(center_x - new_width // 2, 0)
     y1 = max(center_y - new_height // 2, 0)
     x2 = min(center_x + new_width // 2, width)
-    y2 = min(center_y + new_height // 2, height)
+    y2 = min(center_y + new_width // 2, height)
     return image[y1:y2, x1:x2]
 
 def exclude_scaling_box(image: np.ndarray, bbox: np.ndarray, orig_size: tuple, processed_size: tuple, expansion_factor: float = 1.2) -> np.ndarray:
@@ -347,7 +347,7 @@ 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, offset_unit):
+def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width, offset_unit, annotation_text=""):
     msp = doc.modelspace()
     # First, if unit is mm, check if values seem too low (accidental inches) and convert them.
     if offset_unit.lower() == "mm":
@@ -380,24 +380,26 @@ def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width
     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
+    # Set clearance margins.
+    clearance_side = 0.25    # Left and right clearance
+    clearance_top = 0.25     # Top clearance
+    # Bottom clearance: 0.75 if annotation text is provided; otherwise 0.25.
+    clearance_bottom = 0.75 if annotation_text.strip() else 0.25
 
     # 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:
+    if boundary_width_in <= inner_width + 2 * clearance_side or \
+       boundary_length_in <= inner_length + clearance_top + clearance_bottom:
         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
-    half_l = boundary_length_in / 2.0
-    left = shape_cx - half_w
-    right = shape_cx + half_w
-    bottom = shape_cy - half_l
-    top = shape_cy + half_l
-    rect_coords = [(left, bottom), (right, bottom), (right, top), (left, top), (left, bottom)]
+    left = shape_cx - boundary_width_in / 2
+    right = shape_cx + boundary_width_in / 2
+    # Align bottom to inner bounding box with the desired bottom clearance.
+    bottom = min_y - clearance_bottom
+    top = bottom + boundary_length_in
 
+    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"})
@@ -446,11 +448,16 @@ 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 isinstance(image, np.ndarray):
         pil_image = Image.fromarray(image)
         enhanced_image = ImageEnhance.Sharpness(pil_image).enhance(1.5)
         image = np.array(enhanced_image)
+
+    # ---------------------
+    # 1) Detect the drawer with YOLOWorld
+    # ---------------------
     try:
         t = time.time()
         drawer_img = yolo_detect(image)
@@ -462,16 +469,25 @@ def predict(
         print("Image shrinking completed in {:.2f} seconds".format(time.time() - t))
     except DrawerNotDetectedError:
         raise DrawerNotDetectedError("Drawer not detected! Please take another picture with a drawer.")
+
+    # ---------------------
+    # 2) Detect the reference box with YOLO
+    # ---------------------
     try:
         t = time.time()
         reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
         print("Reference square detection completed in {:.2f} seconds".format(time.time() - t))
     except ReferenceBoxNotDetectedError:
         raise ReferenceBoxNotDetectedError("Reference box not detected! Please take another picture with a reference box.")
+
+    # ---------------------
+    # 3) Remove background of the reference box to compute scaling factor
+    # ---------------------
     t = time.time()
     reference_obj_img = make_square(reference_obj_img)
     reference_square_mask = remove_bg_u2netp(reference_obj_img)
     print("Reference image processing completed in {:.2f} seconds".format(time.time() - t))
+
     t = time.time()
     try:
         cv2.imwrite("mask.jpg", cv2.cvtColor(reference_obj_img, cv2.COLOR_RGB2GRAY))
@@ -486,20 +502,21 @@ def predict(
     except Exception as e:
         scaling_factor = None
         print(f"Error calculating scaling factor: {e}")
+
     if scaling_factor is None or scaling_factor == 0:
         scaling_factor = 1.0
         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.
+    # 4) Optional boundary dimension checks
     # ---------------------
     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 units are mm, convert them to inches if the values look small
         if offset_unit.lower() == "mm":
             if boundary_length < 50:
                 boundary_length = boundary_length * 25.4
@@ -510,49 +527,77 @@ def predict(
         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.
+
+        # Basic check so user doesn't request an impossible boundary
+        if boundary_length_in > (image_height_in - 0.5) or boundary_width_in > (image_width_in - 0.5):
+            raise BoundaryExceedsError(
+                "Error: The specified boundary dimensions exceed the allowed image dimensions. Please enter smaller values."
+            )
+
+    # ---------------------
+    # 5) Remove background from the shrunked drawer image (main objects)
+    # ---------------------
     if offset_unit.lower() == "mm":
         if offset_value < 1:
             offset_value = offset_value * 25.4
         offset_inches = offset_value / 25.4
     else:
         offset_inches = offset_value
+
     t = time.time()
     orig_size = shrunked_img.shape[:2]
     objects_mask = remove_bg(shrunked_img)
     processed_size = objects_mask.shape[:2]
+
+    # Exclude the reference box region from the mask
     objects_mask = exclude_scaling_box(objects_mask, scaling_box_coords, orig_size, processed_size, expansion_factor=1.2)
     objects_mask = resize_img(objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0]))
     del scaling_box_coords
     gc.collect()
     print("Object masking completed in {:.2f} seconds".format(time.time() - t))
+
+    # Dilate mask by offset_pixels
     t = time.time()
     offset_pixels = (offset_inches / scaling_factor) * 2 + 1 if scaling_factor != 0 else 1
     dilated_mask = cv2.dilate(objects_mask, np.ones((int(offset_pixels), int(offset_pixels)), np.uint8))
     del objects_mask
     gc.collect()
     print("Mask dilation completed in {:.2f} seconds".format(time.time() - t))
+
     Image.fromarray(dilated_mask).save("./outputs/scaled_mask_new.jpg")
+
+    # ---------------------
+    # 6) Extract outlines from the mask and convert them to DXF splines
+    # ---------------------
     t = time.time()
     outlines, contours = extract_outlines(dilated_mask)
     print("Outline extraction completed in {:.2f} seconds".format(time.time() - t))
+
     output_img = shrunked_img.copy()
     del shrunked_img
     gc.collect()
+
     t = time.time()
     use_finger_clearance = True if finger_clearance.lower() == "yes" else False
-    doc, final_polygons_inch = save_dxf_spline(contours, scaling_factor, processed_size[0], finger_clearance=use_finger_clearance)
+    doc, final_polygons_inch = save_dxf_spline(
+        contours, scaling_factor, processed_size[0], finger_clearance=use_finger_clearance
+    )
     del contours
     gc.collect()
     print("DXF generation completed in {:.2f} seconds".format(time.time() - t))
+
+    # ---------------------
+    # 7) Add optional rectangular boundary (with the 0.75 bottom margin if annotation text is provided)
+    # ---------------------
     boundary_polygon = None
     if add_boundary.lower() == "yes":
-        boundary_polygon = add_rectangular_boundary(doc, final_polygons_inch, boundary_length, boundary_width, offset_unit)
+        boundary_polygon = add_rectangular_boundary(
+            doc, final_polygons_inch, boundary_length, boundary_width, offset_unit, annotation_text
+        )
         if boundary_polygon is not None:
             final_polygons_inch.append(boundary_polygon)
+
+    # Compute bounding box of all polygons to know where to place text
     min_x = float("inf")
     min_y = float("inf")
     max_x = -float("inf")
@@ -567,36 +612,81 @@ def predict(
             max_x = b[2]
         if b[3] > max_y:
             max_y = b[3]
-    margin = 0.40
-    text_x = (min_x + max_x) / 2
-    if add_boundary.lower() == "yes":
-        text_y = min_y + margin
-    else:
-        text_y = min_y - margin
+
+    # ---------------------
+    # 8) Add annotation text (if provided) in the DXF
+    #    - Text is 0.5 inches high
+    #    - Placed so that there is 0.125 inch ABOVE min_y
+    # ---------------------
     msp = doc.modelspace()
     if annotation_text.strip():
+        text_x = (min_x + max_x) / 2.0  # horizontally centered
+        text_height_dxf = 0.5
+
+        # If boundary is "yes", we'll place the baseline right above the min_y (0.125 above)
+        if add_boundary.lower() == "yes":
+            text_y_dxf = min_y + 0.125
+        else:
+            # If no boundary, place it 0.125 above the min_y
+            text_y_dxf = min_y + 0.125
+
         text_entity = msp.add_text(
             annotation_text.strip(),
             dxfattribs={
-                "height": 0.50,
+                "height": text_height_dxf,
                 "layer": "ANNOTATION",
                 "style": "Bold"
             }
         )
-        text_entity.dxf.insert = (text_x, text_y)
+        text_entity.dxf.insert = (text_x, text_y_dxf)
+
+    # Save the DXF
     dxf_filepath = os.path.join("./outputs", "out.dxf")
     doc.saveas(dxf_filepath)
-    draw_polygons_inch(final_polygons_inch, output_img, scaling_factor, processed_size[0], color=(0,0,255), thickness=2)
+
+    # ---------------------
+    # 9) For the preview images, draw the polygons and place text similarly
+    # ---------------------
+    draw_polygons_inch(final_polygons_inch, output_img, scaling_factor, processed_size[0], color=(0, 0, 255), thickness=2)
     new_outlines = np.ones_like(output_img) * 255
-    draw_polygons_inch(final_polygons_inch, new_outlines, scaling_factor, processed_size[0], color=(0,0,255), thickness=2)
+    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))
-        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)
+        text_height_cv = 0.5  # same "logical" height in inches
+        if add_boundary.lower() == "yes":
+            text_y_in = min_y + 0.125
+        else:
+            text_y_in = min_y + 0.125
+
+        # Convert to pixel coords
+        text_px = int(((min_x + max_x) / 2.0) / scaling_factor)
+        text_py = int(processed_size[0] - (text_y_in / scaling_factor))
+        org = (text_px - int(len(annotation_text.strip()) * 6), text_py)  # shift left so text is centered
+
+        cv2.putText(
+            output_img,
+            annotation_text.strip(),
+            org,
+            cv2.FONT_HERSHEY_SIMPLEX,
+            1.3,          # scale factor
+            (0, 0, 255),
+            3,
+            cv2.LINE_AA
+        )
+        cv2.putText(
+            new_outlines,
+            annotation_text.strip(),
+            org,
+            cv2.FONT_HERSHEY_SIMPLEX,
+            1.3,
+            (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 (
         cv2.cvtColor(output_img, cv2.COLOR_BGR2RGB),
         outlines_color,