averile commited on
Commit
9b1d8f1
·
1 Parent(s): 5bf84ff

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +13 -8
app.py CHANGED
@@ -3,21 +3,26 @@ import yolov5
3
  import os
4
  from transformers import pipeline
5
 
6
- ImageClassifier = pipeline(task="image-classification", model="")
7
 
8
  model = yolov5.load('./gentle-meadow.pt', device='cpu')
9
 
10
  def predict(image):
11
- inp = transforms.ToTensor()(inp).unsqueeze(0)
12
- with torch.no_grad():
13
- prediction = torch.nn.functional.softmax(model(inp)[0], dim=0)
14
- confidences = {labels[i]: float(prediction[i]) for i in range(1000)}
15
- return confidences
 
 
 
 
 
 
16
 
17
  demo = gr.Interface(fn=predict,
18
  inputs=gr.inputs.Image(type="pil"),
19
- outputs=gr.outputs.Label(num_top_classes=3),
20
- examples=[["cheetah.jpg"]],
21
  )
22
 
23
  demo.launch()
 
3
  import os
4
  from transformers import pipeline
5
 
6
+ #ImageClassifier = pipeline(task="image-classification", model="")
7
 
8
  model = yolov5.load('./gentle-meadow.pt', device='cpu')
9
 
10
  def predict(image):
11
+ results = model([image], size=224)
12
+ #predictions = imageClassifier(image)
13
+ # classMappings = {
14
+ # 'police': "Police / Authorized Personnel",
15
+ # 'public': 'Unauthorized Person'
16
+ # }
17
+ # output = {}
18
+ # for item in predictions:
19
+ # output[classMappings[item['label']]] = item['score']
20
+
21
+ return results.render()[0]
22
 
23
  demo = gr.Interface(fn=predict,
24
  inputs=gr.inputs.Image(type="pil"),
25
+ outputs=gr.outputs.Image(type="pil"),
 
26
  )
27
 
28
  demo.launch()