Spaces:
Sleeping
Sleeping
File size: 3,250 Bytes
9b658e7 3a4ff58 9b658e7 b0a21e7 3a4ff58 9b658e7 3a4ff58 9b658e7 3a4ff58 9b658e7 2ccf6ca 9b658e7 3a4ff58 b0a21e7 3a4ff58 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
# app.py
import gradio as gr
#import spaces
#import torch
from PIL import Image
from transformers import pipeline
import matplotlib.pyplot as plt
import io
import os
list_models = ["facebook/detr-resnet-50", "facebook/detr-resnet-101", "hustvl/yolos-tiny", "hustvl/yolos-small"]
list_models_simple = [os.path.basename(model) for model in list_models]
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],
]
def load_pipeline(model):
model_pipeline = pipeline(model=model)
return model_pipeline
def get_output_figure(pil_img, results, threshold):
plt.figure(figsize=(16, 10))
plt.imshow(pil_img)
ax = plt.gca()
colors = COLORS * 100
for result in results:
score = result["score"]
label = result["label"]
box = list(result["box"].values())
if score > threshold:
c = COLORS[hash(label) % len(COLORS)]
ax.add_patch(
plt.Rectangle((box[0], box[1]), box[2] - box[0], box[3] - box[1], fill=False, color=c, linewidth=3)
)
text = f"{label}: {score:0.2f}"
ax.text(box[0], box[1], text, fontsize=15, bbox=dict(facecolor="yellow", alpha=0.5))
plt.axis("off")
return plt.gcf()
#@spaces.GPU
def detect(image, model_id, threshold=0.9):
print("model:", list_models[model_id])
model_pipeline = load_pipeline(list_models[model_id])
results = model_pipeline(image)
print(results)
output_figure = get_output_figure(image, results, threshold=threshold)
buf = io.BytesIO()
output_figure.savefig(buf, bbox_inches="tight")
buf.seek(0)
output_pil_img = Image.open(buf)
return output_pil_img
def demo():
with gr.Blocks(theme="base") as demo:
gr.Markdown("# Object detection on COCO dataset")
gr.Markdown(
"""
This application uses transformer-based models to detect objects on images.
This version was trained using the COCO dataset.
You can load an image and see the predictions for the objects detected.
"""
)
with gr.Row():
model_id = gr.Radio(list_models, \
label="Detection models", value=list_models[3], type="index", info="Choose your detection model")
with gr.Row():
threshold = gr.Slider(0, 1.0, value=0.9, label='Detection threshold', info="Choose your detection threshold")
with gr.Row():
input_image = gr.Image(label="Input image", type="pil")
output_image = gr.Image(label="Output image", type="pil")
with gr.Row():
submit_btn = gr.Button("Submit")
clear_button = gr.ClearButton()
gr.Examples(['samples/savanna.jpg'], inputs=input_image)
submit_btn.click(fn=detect, inputs=[input_image, model_id, threshold], outputs=[output_image])
clear_button.click(lambda: [None, None], \
inputs=None, \
outputs=[input_image, output_image], \
queue=False)
demo.queue().launch(debug=True)
if __name__ == "__main__":
demo() |