hlydecker commited on
Commit
5847407
1 Parent(s): 3502d91

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -4
app.py CHANGED
@@ -10,19 +10,21 @@ from PIL import Image
10
  models = ["model_weights/md_v5a.0.0.pt","model_weights/md_v5b.0.0.pt"]
11
  model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/md_v5a.0.0.pt")
12
 
13
- def yolo(im, model):
14
  g = (size / max(im.size)) # gain
15
  im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS) # resize
16
-
 
 
17
  results = model(im) # inference
18
  results.render() # updates results.imgs with boxes and labels
19
  return Image.fromarray(results.imgs[0])
20
 
21
  image = gr.inputs.Image(type="pil", label="Input Image")
22
- model = gr.inputs.Model(choices = models, value = "model_weights/md_v5a.0.0.pt",type = "value", label="Model Weight")
23
  size = 640
24
 
25
- inputs = [image, model, size]
26
 
27
  outputs = gr.outputs.Image(type="pil", label="Output Image")
28
 
 
10
  models = ["model_weights/md_v5a.0.0.pt","model_weights/md_v5b.0.0.pt"]
11
  model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/md_v5a.0.0.pt")
12
 
13
+ def yolo(im, chosen_model, size):
14
  g = (size / max(im.size)) # gain
15
  im = im.resize((int(x * g) for x in im.size), Image.ANTIALIAS) # resize
16
+
17
+ model = torch.hub.load('ultralytics/yolov5', 'custom', "model_weights/md_v5a.0.0.pt")
18
+
19
  results = model(im) # inference
20
  results.render() # updates results.imgs with boxes and labels
21
  return Image.fromarray(results.imgs[0])
22
 
23
  image = gr.inputs.Image(type="pil", label="Input Image")
24
+ chosen_model = gr.inputs.Dropdown(choices = models, value = "model_weights/md_v5a.0.0.pt",type = "value", label="Model Weight")
25
  size = 640
26
 
27
+ inputs = [image, chosen_model, size]
28
 
29
  outputs = gr.outputs.Image(type="pil", label="Output Image")
30