detr / app.py
ayaanzaveri's picture
Create app.py
2387c38
raw
history blame
1.82 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
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)
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')
output = pipe(img_url)
fpath = os.path.join(r"Poppins-SemiBold.ttf")
prop = font_manager.FontProperties(fname=fpath)
img = Image.open('kitchen.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=100, color=selected_color, fontproperties=prop)
plt.axis('off')
plt.savefig('kitchen-bbox.jpg', dpi=1200, bbox_inches='tight')
image = Image.open('kitchen-bbox.jpg')
st.image(image, caption='DETR Image')
plt.show()