JosephTK commited on
Commit
6784e11
1 Parent(s): 8be87bf

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +50 -0
app.py ADDED
@@ -0,0 +1,50 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoImageProcessor, ResNetForImageClassification, YolosFeatureExtractor, YolosForObjectDetection
3
+ import torch
4
+
5
+
6
+
7
+
8
+ def detect(image1, image2):
9
+ ### Image 1, the object ###
10
+ processor = AutoImageProcessor.from_pretrained("microsoft/resnet-50")
11
+ model = ResNetForImageClassification.from_pretrained("microsoft/resnet-50")
12
+
13
+ inputs = processor(image, return_tensors="pt")
14
+
15
+ with torch.no_grad():
16
+ logits = model(**inputs).logits
17
+
18
+ # model predicts one of the 1000 ImageNet classes
19
+ predicted_label = logits.argmax(-1).item()
20
+ print(model.config.id2label[predicted_label])
21
+
22
+
23
+
24
+ ### Image 2, object detections ###
25
+ from PIL import Image
26
+ import requests
27
+
28
+ feature_extractor = YolosFeatureExtractor.from_pretrained('hustvl/yolos-small')
29
+ model = YolosForObjectDetection.from_pretrained('hustvl/yolos-small')
30
+
31
+ inputs = feature_extractor(images=image, return_tensors="pt")
32
+ outputs = model(**inputs)
33
+
34
+ # model predicts bounding boxes and corresponding COCO classes
35
+ logits = outputs.logits
36
+ bboxes = outputs.pred_boxes
37
+
38
+ return model.config.id2label[predicted_label], bboxes
39
+
40
+
41
+
42
+
43
+ demo = gr.Interface(
44
+ fn=detect,
45
+ inputs=[gr.inputs.Image(label="Object to detect"), gr.inputs.Image(label="Image to detect object in")],
46
+ outputs=["prediction", "bounding boxes"],
47
+ title="Object Counts in Image"
48
+ )
49
+
50
+ demo.launch()