ammariii08 commited on
Commit
54bcad2
·
verified ·
1 Parent(s): 62238e9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +17 -25
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
- return None, None, None, None, f"Error: {str(e)}"
473
- # Ensure that shrunked_img is defined only after successful detection
 
 
 
474
  t = time.time()
475
- shrunked_img = make_square(shrink_bbox(drawer_img, 0.90))
 
 
 
 
 
476
  del drawer_img
477
  gc.collect()
478
- print("Image shrinking completed in {:.2f} seconds".format(time.time() - t))
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
  # ---------------------