NimaBoscarino
commited on
Commit
•
27d3848
1
Parent(s):
aab5e95
Update app.py
Browse files
app.py
CHANGED
@@ -20,6 +20,7 @@ from PIL import Image
|
|
20 |
|
21 |
from huggingface_hub import hf_hub_download
|
22 |
|
|
|
23 |
def load_model(model_name):
|
24 |
model_path = hf_hub_download(repo_id=f"Yolov7/{model_name}", filename=f"{model_name}.pt")
|
25 |
|
@@ -32,10 +33,12 @@ model_names = [
|
|
32 |
"yolov7-e6",
|
33 |
]
|
34 |
|
|
|
|
|
35 |
|
36 |
def detect(img,model):
|
37 |
parser = argparse.ArgumentParser()
|
38 |
-
parser.add_argument('--weights', nargs='+', type=str, default=
|
39 |
parser.add_argument('--source', type=str, default='Inference/', help='source') # file/folder, 0 for webcam
|
40 |
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
41 |
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
|
|
|
20 |
|
21 |
from huggingface_hub import hf_hub_download
|
22 |
|
23 |
+
|
24 |
def load_model(model_name):
|
25 |
model_path = hf_hub_download(repo_id=f"Yolov7/{model_name}", filename=f"{model_name}.pt")
|
26 |
|
|
|
33 |
"yolov7-e6",
|
34 |
]
|
35 |
|
36 |
+
models = {model_name: load_model(model_name) for model_name in model_names}
|
37 |
+
|
38 |
|
39 |
def detect(img,model):
|
40 |
parser = argparse.ArgumentParser()
|
41 |
+
parser.add_argument('--weights', nargs='+', type=str, default=models[model], help='model.pt path(s)')
|
42 |
parser.add_argument('--source', type=str, default='Inference/', help='source') # file/folder, 0 for webcam
|
43 |
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
|
44 |
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
|