rachana219 commited on
Commit
2225fd6
·
1 Parent(s): 1c39508

fixing issue with image detection

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -138,7 +138,7 @@ def inference(img,model_link,iou_threshold,confidence_threshold):
138
  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
139
  # Load model
140
  model_path = 'weights/'+str(model_link)+'.pt'
141
- if model_link== '.':
142
  model = YOLO(model_path)
143
  return detectv8(img,model,device,iou_threshold,confidence_threshold)
144
  else:
 
138
  device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
139
  # Load model
140
  model_path = 'weights/'+str(model_link)+'.pt'
141
+ if model_link== 'yolov8m':
142
  model = YOLO(model_path)
143
  return detectv8(img,model,device,iou_threshold,confidence_threshold)
144
  else: