Update app.py

app.py

@@ -36,23 +36,6 @@ transform_image = transforms.Compose(
 )
 
 
-def yolo_detect(
-    image: Union[str, Path, int, Image.Image, list, tuple, np.ndarray, torch.Tensor],
-    classes: List[str],
-) -> np.ndarray:
-    drawer_detector = YOLOWorld("yolov8x-worldv2.pt")
-    drawer_detector.set_classes(classes)
-    results: List[Results] = drawer_detector.predict(image)
-    boxes = []
-    for result in results:
-        boxes.append(
-            save_one_box(result.cpu().boxes.xyxy, im=result.orig_img, save=False)
-        )
-
-    del drawer_detector
-
-    return boxes[0]
-
 
 def remove_bg(image: np.ndarray) -> np.ndarray:
     image = Image.fromarray(image)
@@ -338,17 +321,12 @@ def resize_img(img: np.ndarray, resize_dim):
 
 
 def predict(image, offset_inches):
-    try:
-        drawer_img = yolo_detect(image, ["box"])
-        shrunked_img = make_square(shrink_bbox(drawer_img, 0.8))
-    except:
-        raise gr.Error("Unable to DETECT DRAWER, please take another picture with different magnification level!")
-
+    
     # Detect the scaling reference square
     try:
-        reference_obj_img, scaling_box_coords = detect_reference_square(shrunked_img)
+        reference_obj_img, scaling_box_coords = detect_reference_square(image)
     except:
-        raise gr.Error("Unable to DETECT REFERENCE BOX, please take another picture with different magnification level!")
+        raise gr.Error("Unable to DETECT COIN, please take another picture with different magnification level!")
 
     # reference_obj_img_scaled = shrink_bbox(reference_obj_img, 1.2)
     # make the image sqaure so it does not effect the size of objects
@@ -370,9 +348,9 @@ def predict(image, offset_inches):
         scaling_factor = 1.0
 
     # Save original size before `remove_bg` processing
-    orig_size = shrunked_img.shape[:2]
+    orig_size = image.shape[:2]
     # Generate foreground mask and save its size
-    objects_mask = remove_bg(shrunked_img)
+    objects_mask = remove_bg(image)
 
     processed_size = objects_mask.shape[:2]
     # Exclude scaling box region from objects mask
@@ -384,7 +362,7 @@ def predict(image, offset_inches):
         expansion_factor=3.0,
     )
     objects_mask = resize_img(
-        objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0])
+        objects_mask, (image.shape[1], image.shape[0])
     )
     offset_pixels = (offset_inches / scaling_factor) * 2 + 1
     dilated_mask = cv2.dilate(