Spaces:
Sleeping
Sleeping
Update main.py
Browse files
main.py
CHANGED
|
@@ -32,9 +32,9 @@ from ultralytics import YOLO
|
|
| 32 |
# os.makedirs("/data/icon_detect")
|
| 33 |
|
| 34 |
try:
|
| 35 |
-
yolo_model = YOLO("weights/
|
| 36 |
except:
|
| 37 |
-
yolo_model = YOLO("weights/
|
| 38 |
|
| 39 |
from transformers import AutoProcessor, AutoModelForCausalLM
|
| 40 |
|
|
|
|
| 32 |
# os.makedirs("/data/icon_detect")
|
| 33 |
|
| 34 |
try:
|
| 35 |
+
yolo_model = YOLO("weights/best.pt").to("cuda")
|
| 36 |
except:
|
| 37 |
+
yolo_model = YOLO("weights/best.pt")
|
| 38 |
|
| 39 |
from transformers import AutoProcessor, AutoModelForCausalLM
|
| 40 |
|