Nighty3912 commited on
Commit
18f47a4
·
verified ·
1 Parent(s): a0e9cd0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -72,10 +72,10 @@ print("YOLOWorld model loaded in {:.2f} seconds".format(time.time() - start_time
72
 
73
  print("Loading YOLO reference model...")
74
  start_time = time.time()
75
- reference_model_path = os.path.join(CACHE_DIR, "coin_det.pt")
76
  if not os.path.exists(reference_model_path):
77
  print("Caching YOLO reference model to", reference_model_path)
78
- shutil.copy("coin_det.pt", reference_model_path)
79
  reference_detector_global = YOLO(reference_model_path)
80
  print("YOLO reference model loaded in {:.2f} seconds".format(time.time() - start_time))
81
 
@@ -123,7 +123,7 @@ def unload_and_reload_models():
123
  gc.collect()
124
  new_drawer_detector = YOLOWorld(os.path.join(CACHE_DIR, "yolov8x-worldv2.pt"))
125
  new_drawer_detector.set_classes(["box"])
126
- new_reference_detector = YOLO(os.path.join(CACHE_DIR, "coin_det.pt"))
127
  new_birefnet = AutoModelForImageSegmentation.from_pretrained(
128
  "zhengpeng7/BiRefNet", trust_remote_code=True, cache_dir=CACHE_DIR
129
  )
@@ -158,7 +158,7 @@ def yolo_detect(image: Union[str, Path, int, Image.Image, list, tuple, np.ndarra
158
 
159
  def detect_reference_square(img: np.ndarray):
160
  t = time.time()
161
- res = reference_detector_global.predict(img, conf=0.15)
162
  if not res or len(res) == 0 or len(res[0].boxes) == 0:
163
  raise ReferenceBoxNotDetectedError("Reference Coin not detected in the image.")
164
  print("Reference coin detection completed in {:.2f} seconds".format(time.time() - t))
 
72
 
73
  print("Loading YOLO reference model...")
74
  start_time = time.time()
75
+ reference_model_path = os.path.join(CACHE_DIR, "best_coin.pt")
76
  if not os.path.exists(reference_model_path):
77
  print("Caching YOLO reference model to", reference_model_path)
78
+ shutil.copy("best_coin.pt", reference_model_path)
79
  reference_detector_global = YOLO(reference_model_path)
80
  print("YOLO reference model loaded in {:.2f} seconds".format(time.time() - start_time))
81
 
 
123
  gc.collect()
124
  new_drawer_detector = YOLOWorld(os.path.join(CACHE_DIR, "yolov8x-worldv2.pt"))
125
  new_drawer_detector.set_classes(["box"])
126
+ new_reference_detector = YOLO(os.path.join(CACHE_DIR, "best_coin.pt"))
127
  new_birefnet = AutoModelForImageSegmentation.from_pretrained(
128
  "zhengpeng7/BiRefNet", trust_remote_code=True, cache_dir=CACHE_DIR
129
  )
 
158
 
159
  def detect_reference_square(img: np.ndarray):
160
  t = time.time()
161
+ res = reference_detector_global.predict(img, conf=0.5)
162
  if not res or len(res) == 0 or len(res[0].boxes) == 0:
163
  raise ReferenceBoxNotDetectedError("Reference Coin not detected in the image.")
164
  print("Reference coin detection completed in {:.2f} seconds".format(time.time() - t))