ailm commited on
Commit
aa09382
1 Parent(s): 22f5e80

update app.py

Browse files
Files changed (1) hide show
  1. app.py +65 -0
app.py CHANGED
@@ -0,0 +1,65 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import cv2
2
+ import numpy as np
3
+ import gradio as gr
4
+ import matplotlib.pyplot as plt
5
+ from transformers import pipeline
6
+
7
+
8
+ model = pipeline("object-detection", "facebook/detr-resnet-50") #loading model
9
+
10
+ #render function
11
+
12
+ def render_results(raw_image, model_output):
13
+
14
+ raw_image = np.array(raw_image)
15
+
16
+ for detection in model_output:
17
+ label = detection['label']
18
+ score = detection['score']
19
+ box = detection['box']
20
+ xmin, ymin, xmax, ymax = box['xmin'], box['ymin'], box['xmax'], box['ymax']
21
+
22
+ #Drawing the bounding box
23
+ cv2.rectangle(raw_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
24
+
25
+ #Puting label and score near the bounding box
26
+ cv2.putText(raw_image, f"{label}: {score:.2f}", (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
27
+ return raw_image
28
+
29
+
30
+
31
+ def get_object_counts(detections): ##to get count of object detected in the image
32
+ object_counts = {}
33
+ for detection in detections:
34
+ label = detection['label']
35
+ if label in object_counts:
36
+ object_counts[label] += 1
37
+ else:
38
+ object_counts[label] = 1
39
+
40
+ return object_counts
41
+
42
+ def generate_output_text(object_counts): ##to get the output string
43
+ output_text = "In this image there are"
44
+ for label, count in object_counts.items():
45
+ output_text += f" {count} {label},"
46
+ output_text = output_text.rstrip(',') + "."
47
+ return output_text
48
+
49
+
50
+
51
+ def main(pil_image):
52
+ pipeline_output = model(pil_image) #model output
53
+ processed_image = render_results(pil_image, pipeline_output) ##process image by drawing bounding boxes
54
+ output_text = generate_output_text(get_object_counts(pipeline_output)) ##output string
55
+
56
+ return processed_image, output_text
57
+
58
+
59
+ demo = gr.Interface(
60
+ fn = main,
61
+ inputs = gr.Image(label = "Input Image", type = "pil"),
62
+ outputs = [gr.Image(label = "Modle output Predictions", type = "numpy"), gr.Text(label="Output Text")]
63
+ )
64
+
65
+ demo.launch()