Spaces:
Sleeping
Sleeping
Update main.py
Browse files
main.py
CHANGED
@@ -23,18 +23,18 @@ from utils import (
|
|
23 |
)
|
24 |
import torch
|
25 |
|
26 |
-
|
27 |
-
|
28 |
|
29 |
from ultralytics import YOLO
|
30 |
|
31 |
-
|
32 |
-
|
33 |
|
34 |
try:
|
35 |
-
yolo_model = YOLO("best.pt").to("cuda")
|
36 |
except:
|
37 |
-
yolo_model = YOLO("best.pt")
|
38 |
|
39 |
from transformers import AutoProcessor, AutoModelForCausalLM
|
40 |
|
@@ -129,4 +129,4 @@ async def process_image(
|
|
129 |
raise HTTPException(status_code=400, detail="Invalid image file")
|
130 |
|
131 |
response = process(image_input, box_threshold, iou_threshold)
|
132 |
-
return response
|
|
|
23 |
)
|
24 |
import torch
|
25 |
|
26 |
+
yolo_model = get_yolo_model(model_path='/data/icon_detect/best.pt')
|
27 |
+
caption_model_processor = get_caption_model_processor(model_name="florence2", model_name_or_path="/data/icon_caption_florence")
|
28 |
|
29 |
from ultralytics import YOLO
|
30 |
|
31 |
+
if not os.path.exists("/data/icon_detect"):
|
32 |
+
os.makedirs("/data/icon_detect")
|
33 |
|
34 |
try:
|
35 |
+
yolo_model = YOLO("OmniPar/best.pt").to("cuda")
|
36 |
except:
|
37 |
+
yolo_model = YOLO("OmniPar/best.pt")
|
38 |
|
39 |
from transformers import AutoProcessor, AutoModelForCausalLM
|
40 |
|
|
|
129 |
raise HTTPException(status_code=400, detail="Invalid image file")
|
130 |
|
131 |
response = process(image_input, box_threshold, iou_threshold)
|
132 |
+
return response
|