detr / app.py
ayaanzaveri's picture
Update app.py
6452d66
raw
history blame
2.15 kB
from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from random import choice
from PIL import Image
import os
from matplotlib import rcParams, font_manager
import streamlit as st
import urllib.request
import requests
extractor = AutoFeatureExtractor.from_pretrained("facebook/detr-resnet-50")
model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-50")
from transformers import pipeline
pipe = pipeline('object-detection', model=model, feature_extractor=extractor)
my_bar = st.progress(0)
img_url = st.text_input('Image URL', 'https://images.unsplash.com/photo-1556911220-bff31c812dba?ixlib=rb-4.0.3&ixid=MnwxMjA3fDB8MHxwaG90by1wYWdlfHx8fGVufDB8fHx8&auto=format&fit=crop&w=2468&q=80')
my_bar.progress(percent_complete + 20)
img_data = requests.get(img_url).content
with open('detect.jpg', 'wb') as handler:
handler.write(img_data)
my_bar.progress(percent_complete + 40)
output = pipe(img_url)
my_bar.progress(percent_complete + 20)
fpath = "Poppins-SemiBold.ttf"
prop = font_manager.FontProperties(fname=fpath)
img = Image.open('detect.jpg')
plt.figure(dpi=2400)
# Create figure and axes
fig, ax = plt.subplots()
# Display the image
ax.imshow(img)
colors = ["#ef4444", "#f97316", "#eab308", "#84cc16", "#06b6d4", "#6366f1"]
# Create a Rectangle patch
for prediction in output:
selected_color = choice(colors)
x, y, w, h = prediction['box']['xmin'], prediction['box']['ymin'], prediction['box']['xmax'] - prediction['box']['xmin'], prediction['box']['ymax'] - prediction['box']['ymin']
rect = patches.FancyBboxPatch((x, y), w, h, linewidth=1.25, edgecolor=selected_color, facecolor='none', boxstyle="round,pad=-0.0040,rounding_size=10",)
ax.add_patch(rect)
plt.text(x, y-25, f"{prediction['label']}: {round(prediction['score']*100, 1)}%", fontsize=5, color=selected_color, fontproperties=prop)
plt.axis('off')
plt.savefig('detect-bbox.jpg', dpi=1200, bbox_inches='tight')
image = Image.open('detect-bbox.jpg')
st.image(image, caption='DETR Image')
plt.show()
my_bar.progress(percent_complete + 20)