Spaces:
Running
Running
Update app.py
Browse files
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, "
|
76 |
if not os.path.exists(reference_model_path):
|
77 |
print("Caching YOLO reference model to", reference_model_path)
|
78 |
-
shutil.copy("
|
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, "
|
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.
|
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))
|