Update app.py

app.py

@@ -458,7 +458,7 @@ def predict(
     # Apply brightness and sharpness enhancement.
     if isinstance(image, np.ndarray):
         pil_image = Image.fromarray(image)
-        Bright = ImageEnhance.Brightness(pil_image).enhance(1) 
+        Bright = ImageEnhance.Brightness(pil_image).enhance(1)
         enhanced_image = ImageEnhance.Sharpness(Bright).enhance(0.5)
         image = np.array(enhanced_image)
 
@@ -469,13 +469,14 @@ def predict(
         t = time.time()
         drawer_img = yolo_detect(image)
         print("Drawer detection completed in {:.2f} seconds".format(time.time() - t))
-        t = time.time()
-        shrunked_img = make_square(shrink_bbox(drawer_img, 0.90))
-        del drawer_img
-        gc.collect()
-        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.")
+    except DrawerNotDetectedError as e:
+        return None, None, None, None, f"Error: {str(e)}"
+    # Ensure that shrunked_img is defined only after successful detection.
+    t = time.time()
+    shrunked_img = make_square(shrink_bbox(drawer_img, 0.90))
+    del drawer_img
+    gc.collect()
+    print("Image shrinking completed in {:.2f} seconds".format(time.time() - t))
 
     # ---------------------
     # 2) Detect the reference box with YOLO
@@ -484,8 +485,8 @@ def predict(
         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.")
+    except ReferenceBoxNotDetectedError as e:
+        return None, None, None, None, f"Error: {str(e)}"
 
     # ---------------------
     # 3) Remove background of the reference box to compute scaling factor
@@ -523,7 +524,6 @@ def predict(
         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
-        # 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
@@ -535,7 +535,6 @@ def predict(
             boundary_length_in = boundary_length
             boundary_width_in = boundary_width
 
-        # Basic check so user doesn't request an impossible boundary
         if boundary_length_in > (image_height_in - 1) or boundary_width_in > (image_width_in - 1):
             raise BoundaryExceedsError(
                 "Error: The specified boundary dimensions exceed the allowed image dimensions. Please enter smaller values."
@@ -556,7 +555,6 @@ def predict(
     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
@@ -608,7 +606,7 @@ def predict(
         inner_max_y = max(inner_max_y, b[3])
 
     # ---------------------
-    # 7) Add optional rectangular boundary (with the 0.75 bottom margin if annotation text is provided)
+    # 7) Add optional rectangular boundary
     # ---------------------
     boundary_polygon = None
     if add_boundary.lower() == "yes":
@@ -619,16 +617,14 @@ def predict(
             boundary_width,
             offset_unit,
             annotation_text,
-            image_height_in=shrunked_img.shape[0] * scaling_factor,
-            image_width_in=shrunked_img.shape[1] * scaling_factor
+            image_height_in=output_img.shape[0] * scaling_factor,
+            image_width_in=output_img.shape[1] * scaling_factor
         )
         if boundary_polygon is not None:
             final_polygons_inch.append(boundary_polygon)
 
     # ---------------------
     # 8) Add annotation text (if provided) in the DXF
-    #    - Text is 0.5 inches high
-    #    - Placed so that its top edge is 0.125 inches below the inner tools' contour
     # ---------------------
     msp = doc.modelspace()
     if annotation_text.strip():