kadirnar commited on
Commit
7df1d09
·
verified ·
1 Parent(s): 29ee754

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -1
app.py CHANGED
@@ -49,7 +49,11 @@ def attempt_download_from_hub(repo_id, hf_token=None):
49
  @spaces.GPU(duration=10)
50
  def LeYOLO_inference(image, model_id, image_size, conf_threshold, iou_threshold):
51
  # MODEL_PATH = attempt_download_from_hub(model_id)
52
- model = model = YOLO.from_pretrained(model_id)
 
 
 
 
53
  model.to('cuda')
54
  results = model(source=image, imgsz=image_size, iou=iou_threshold, conf=conf_threshold, verbose=False)[0]
55
  detections = sv.Detections.from_ultralytics(results)
 
49
  @spaces.GPU(duration=10)
50
  def LeYOLO_inference(image, model_id, image_size, conf_threshold, iou_threshold):
51
  # MODEL_PATH = attempt_download_from_hub(model_id)
52
+ split_path = model.split('/')
53
+ if len(split_path) == 2 and (not os.path.exists(model)):
54
+ model = YOLO.from_pretrained(model_id)
55
+
56
+
57
  model.to('cuda')
58
  results = model(source=image, imgsz=image_size, iou=iou_threshold, conf=conf_threshold, verbose=False)[0]
59
  detections = sv.Detections.from_ultralytics(results)