Update app.py

app.py

@@ -1,81 +1,66 @@
-import io
-
+from transformers import AutoFeatureExtractor, AutoModelForObjectDetection
 import matplotlib.pyplot as plt
-import requests
-import streamlit as st
-import torch
+import matplotlib.patches as patches
+from random import choice
 from PIL import Image
-from transformers import DetrFeatureExtractor, DetrForObjectDetection
-
-# colors for visualization
-COLORS = [
-    [0.000, 0.447, 0.741],
-    [0.850, 0.325, 0.098],
-    [0.929, 0.694, 0.125],
-    [0.494, 0.184, 0.556],
-    [0.466, 0.674, 0.188],
-    [0.301, 0.745, 0.933]
-]
-
-
-@st.cache(allow_output_mutation=True)
-def get_hf_components(model_name_or_path):
-    feature_extractor = DetrFeatureExtractor.from_pretrained(model_name_or_path)
-    model = DetrForObjectDetection.from_pretrained(model_name_or_path)
-    model.eval()
-    return feature_extractor, model
-
-
-@st.cache
-def get_img_from_url(url):
-    return Image.open(requests.get(url, stream=True).raw)
-
-
-def fig2img(fig):
-    buf = io.BytesIO()
-    fig.savefig(buf)
-    buf.seek(0)
-    img = Image.open(buf)
-    return img
-
-
-def visualize_prediction(pil_img, output_dict, threshold=0.7, id2label=None):
-    keep = output_dict["scores"] > threshold
-    boxes = output_dict["boxes"][keep].tolist()
-    scores = output_dict["scores"][keep].tolist()
-    labels = output_dict["labels"][keep].tolist()
-    if id2label is not None:
-        labels = [id2label[x] for x in labels]
-
-    plt.figure(figsize=(16, 10))
-    plt.imshow(pil_img)
-    ax = plt.gca()
-    colors = COLORS * 100
-    for score, (xmin, ymin, xmax, ymax), label, color in zip(scores, boxes, labels, colors):
-        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=color, linewidth=3))
-        ax.text(xmin, ymin, f"{label}: {score:0.2f}", fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
-    plt.axis("off")
-    return fig2img(plt.gcf())
-
-
-def make_prediction(img, feature_extractor, model):
-    inputs = feature_extractor(img, return_tensors="pt")
-    outputs = model(**inputs)
-    img_size = torch.tensor([tuple(reversed(img.size))])
-    processed_outputs = feature_extractor.post_process(outputs, img_size)
-    return processed_outputs[0]
-
-
-def main():
-    option = st.selectbox("Which model should we use?", ("facebook/detr-resnet-50", "facebook/detr-resnet-101"))
-    feature_extractor, model = get_hf_components(option)
-    url = st.text_input("URL to some image", "http://images.cocodataset.org/val2017/000000039769.jpg")
-    img = get_img_from_url(url)
-    processed_outputs = make_prediction(img, feature_extractor, model)
-    threshold = st.slider("Prediction Threshold", 0.0, 1.0, 0.7)
-    viz_img = visualize_prediction(img, processed_outputs, threshold, model.config.id2label)
-    st.image(viz_img)
-
-
-if __name__ == "__main__":
-    main()
+import os
+from matplotlib import rcParams, font_manager
+import streamlit as st
+import urllib.request
+import requests
+
+extractor = AutoFeatureExtractor.from_pretrained("facebook/detr-resnet-101")
+
+model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-101")
+
+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')
+
+st.caption('Downloading Image...')
+
+img_data = requests.get(img_url).content
+with open('detect.jpg', 'wb') as handler:
+    handler.write(img_data)
+
+st.caption('Running Detection...')
+
+output = pipe(img_url)
+
+st.caption('Adding Predictions to Image...')
+
+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()
+
+st.caption('Done!')