Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -41,10 +41,6 @@ class ReferenceBoxNotDetectedError(Exception):
|
|
41 |
"""Raised when the reference box cannot be detected in the image"""
|
42 |
pass
|
43 |
|
44 |
-
class BoundaryExceedsError(Exception):
|
45 |
-
"""Raised when the optional boundary dimensions exceed allowed image dimensions."""
|
46 |
-
pass
|
47 |
-
|
48 |
class BoundaryOverlapError(Exception):
|
49 |
"""Raised when the optional boundary dimensions are too small and overlap with the inner contours."""
|
50 |
pass
|
@@ -390,11 +386,6 @@ def add_rectangular_boundary(doc, polygons_inch, boundary_length, boundary_width
|
|
390 |
if boundary_width_in <= inner_width + 2 * clearance_side or boundary_length_in <= inner_length + 2 * clearance_tb:
|
391 |
raise BoundaryOverlapError("Error: The specified boundary dimensions are too small and overlap with the inner contours. Please provide larger values.")
|
392 |
|
393 |
-
# Check if boundary exceeds image limits (1 inch less than image dimensions)
|
394 |
-
if image_height_in is not None and image_width_in is not None:
|
395 |
-
if boundary_length_in > (image_height_in - 2) or boundary_width_in > (image_width_in - 2):
|
396 |
-
raise BoundaryExceedsError("Error: The specified boundary dimensions exceed the allowed image dimensions. Please enter smaller values.")
|
397 |
-
|
398 |
# Calculate center of inner contours
|
399 |
center_x = (min_x + max_x) / 2
|
400 |
center_y = (min_y + max_y) / 2
|
@@ -462,23 +453,32 @@ def predict(
|
|
462 |
image = np.array(enhanced_image)
|
463 |
|
464 |
# ---------------------
|
465 |
-
# 1) Detect the drawer with YOLOWorld
|
466 |
# ---------------------
|
|
|
467 |
try:
|
468 |
t = time.time()
|
469 |
drawer_img = yolo_detect(image)
|
470 |
print("Drawer detection completed in {:.2f} seconds".format(time.time() - t))
|
471 |
except DrawerNotDetectedError as e:
|
472 |
-
|
473 |
-
|
|
|
|
|
|
|
474 |
t = time.time()
|
475 |
-
|
|
|
|
|
|
|
|
|
|
|
476 |
del drawer_img
|
477 |
gc.collect()
|
478 |
-
print("Image
|
479 |
|
480 |
# ---------------------
|
481 |
-
# 2) Detect the reference box with YOLO
|
482 |
# ---------------------
|
483 |
try:
|
484 |
t = time.time()
|
@@ -517,12 +517,9 @@ def predict(
|
|
517 |
print("Scaling factor determined: {}".format(scaling_factor))
|
518 |
|
519 |
# ---------------------
|
520 |
-
# 4) Optional boundary dimension checks
|
521 |
# ---------------------
|
522 |
if add_boundary.lower() == "yes":
|
523 |
-
image_height_px, image_width_px = shrunked_img.shape[:2]
|
524 |
-
image_height_in = image_height_px * scaling_factor
|
525 |
-
image_width_in = image_width_px * scaling_factor
|
526 |
if offset_unit.lower() == "mm":
|
527 |
if boundary_length < 50:
|
528 |
boundary_length = boundary_length * 25.4
|
@@ -533,12 +530,7 @@ def predict(
|
|
533 |
else:
|
534 |
boundary_length_in = boundary_length
|
535 |
boundary_width_in = boundary_width
|
536 |
-
|
537 |
-
if boundary_length_in > (image_height_in - 1) or boundary_width_in > (image_width_in - 1):
|
538 |
-
raise BoundaryExceedsError(
|
539 |
-
"Error: The specified boundary dimensions exceed the allowed image dimensions. Please enter smaller values."
|
540 |
-
)
|
541 |
-
|
542 |
# ---------------------
|
543 |
# 5) Remove background from the shrunked drawer image (main objects)
|
544 |
# ---------------------
|
|
|
41 |
"""Raised when the reference box cannot be detected in the image"""
|
42 |
pass
|
43 |
|
|
|
|
|
|
|
|
|
44 |
class BoundaryOverlapError(Exception):
|
45 |
"""Raised when the optional boundary dimensions are too small and overlap with the inner contours."""
|
46 |
pass
|
|
|
386 |
if boundary_width_in <= inner_width + 2 * clearance_side or boundary_length_in <= inner_length + 2 * clearance_tb:
|
387 |
raise BoundaryOverlapError("Error: The specified boundary dimensions are too small and overlap with the inner contours. Please provide larger values.")
|
388 |
|
|
|
|
|
|
|
|
|
|
|
389 |
# Calculate center of inner contours
|
390 |
center_x = (min_x + max_x) / 2
|
391 |
center_y = (min_y + max_y) / 2
|
|
|
453 |
image = np.array(enhanced_image)
|
454 |
|
455 |
# ---------------------
|
456 |
+
# 1) Detect the drawer with YOLOWorld (or use original image if not detected)
|
457 |
# ---------------------
|
458 |
+
drawer_detected = True
|
459 |
try:
|
460 |
t = time.time()
|
461 |
drawer_img = yolo_detect(image)
|
462 |
print("Drawer detection completed in {:.2f} seconds".format(time.time() - t))
|
463 |
except DrawerNotDetectedError as e:
|
464 |
+
print(f"Drawer not detected: {e}, using original image.")
|
465 |
+
drawer_detected = False
|
466 |
+
drawer_img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
467 |
+
|
468 |
+
# Process the image (either cropped drawer or original)
|
469 |
t = time.time()
|
470 |
+
if drawer_detected:
|
471 |
+
# For detected drawers: shrink and square
|
472 |
+
shrunked_img = make_square(shrink_bbox(drawer_img, 0.90))
|
473 |
+
else:
|
474 |
+
# For non-drawer images: keep original dimensions
|
475 |
+
shrunked_img = drawer_img # Already in BGR format from above
|
476 |
del drawer_img
|
477 |
gc.collect()
|
478 |
+
print("Image processing completed in {:.2f} seconds".format(time.time() - t))
|
479 |
|
480 |
# ---------------------
|
481 |
+
# 2) Detect the reference box with YOLO (now works on either cropped or original image)
|
482 |
# ---------------------
|
483 |
try:
|
484 |
t = time.time()
|
|
|
517 |
print("Scaling factor determined: {}".format(scaling_factor))
|
518 |
|
519 |
# ---------------------
|
520 |
+
# 4) Optional boundary dimension checks (now without size limits)
|
521 |
# ---------------------
|
522 |
if add_boundary.lower() == "yes":
|
|
|
|
|
|
|
523 |
if offset_unit.lower() == "mm":
|
524 |
if boundary_length < 50:
|
525 |
boundary_length = boundary_length * 25.4
|
|
|
530 |
else:
|
531 |
boundary_length_in = boundary_length
|
532 |
boundary_width_in = boundary_width
|
533 |
+
|
|
|
|
|
|
|
|
|
|
|
534 |
# ---------------------
|
535 |
# 5) Remove background from the shrunked drawer image (main objects)
|
536 |
# ---------------------
|