Update app.py

app.py

@@ -69,10 +69,10 @@ print("YOLOWorld model loaded in {:.2f} seconds".format(time.time() - start_time
 
 print("Loading YOLO reference model...")
 start_time = time.time()
-reference_model_path = os.path.join(CACHE_DIR, "last.pt")
+reference_model_path = os.path.join(CACHE_DIR, "best.pt")
 if not os.path.exists(reference_model_path):
     print("Caching YOLO reference model to", reference_model_path)
-    shutil.copy("last.pt", reference_model_path)
+    shutil.copy("best.pt", reference_model_path)
 reference_detector_global = YOLO(reference_model_path)
 print("YOLO reference model loaded in {:.2f} seconds".format(time.time() - start_time))
 
@@ -120,7 +120,7 @@ def unload_and_reload_models():
     gc.collect()
     new_drawer_detector = YOLOWorld(os.path.join(CACHE_DIR, "yolov8x-worldv2.pt"))
     new_drawer_detector.set_classes(["box"])
-    new_reference_detector = YOLO(os.path.join(CACHE_DIR, "last.pt"))
+    new_reference_detector = YOLO(os.path.join(CACHE_DIR, "best.pt"))
     new_birefnet = AutoModelForImageSegmentation.from_pretrained(
         "zhengpeng7/BiRefNet", trust_remote_code=True, cache_dir=CACHE_DIR
     )
@@ -155,7 +155,7 @@ def yolo_detect(image: Union[str, Path, int, Image.Image, list, tuple, np.ndarra
 
 def detect_reference_square(img: np.ndarray):
     t = time.time()
-    res = reference_detector_global.predict(img, conf=0.05)
+    res = reference_detector_global.predict(img, conf=0.25)
     if not res or len(res) == 0 or len(res[0].boxes) == 0:
         raise ReferenceBoxNotDetectedError("Reference box not detected in the image.")
     print("Reference detection completed in {:.2f} seconds".format(time.time() - t))
@@ -561,7 +561,7 @@ def predict(
     objects_mask = remove_bg(shrunked_img)
     processed_size = objects_mask.shape[:2]
 
-    objects_mask = exclude_scaling_box(objects_mask, scaling_box_coords, orig_size, processed_size, expansion_factor=1.2)
+    objects_mask = exclude_scaling_box(objects_mask, scaling_box_coords, orig_size, processed_size, expansion_factor=1.7)
     objects_mask = resize_img(objects_mask, (shrunked_img.shape[1], shrunked_img.shape[0]))
     del scaling_box_coords
     gc.collect()