NimaBoscarino commited on
Commit
27d3848
β€’
1 Parent(s): aab5e95

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -1
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=load_model(model), help='model.pt path(s)')
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')