ailm's picture
update app.py
aa09382 verified
raw
history blame contribute delete
No virus
1.94 kB
import cv2
import numpy as np
import gradio as gr
import matplotlib.pyplot as plt
from transformers import pipeline
model = pipeline("object-detection", "facebook/detr-resnet-50") #loading model
#render function
def render_results(raw_image, model_output):
raw_image = np.array(raw_image)
for detection in model_output:
label = detection['label']
score = detection['score']
box = detection['box']
xmin, ymin, xmax, ymax = box['xmin'], box['ymin'], box['xmax'], box['ymax']
#Drawing the bounding box
cv2.rectangle(raw_image, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
#Puting label and score near the bounding box
cv2.putText(raw_image, f"{label}: {score:.2f}", (xmin, ymin - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
return raw_image
def get_object_counts(detections): ##to get count of object detected in the image
object_counts = {}
for detection in detections:
label = detection['label']
if label in object_counts:
object_counts[label] += 1
else:
object_counts[label] = 1
return object_counts
def generate_output_text(object_counts): ##to get the output string
output_text = "In this image there are"
for label, count in object_counts.items():
output_text += f" {count} {label},"
output_text = output_text.rstrip(',') + "."
return output_text
def main(pil_image):
pipeline_output = model(pil_image) #model output
processed_image = render_results(pil_image, pipeline_output) ##process image by drawing bounding boxes
output_text = generate_output_text(get_object_counts(pipeline_output)) ##output string
return processed_image, output_text
demo = gr.Interface(
fn = main,
inputs = gr.Image(label = "Input Image", type = "pil"),
outputs = [gr.Image(label = "Modle output Predictions", type = "numpy"), gr.Text(label="Output Text")]
)
demo.launch()