|
import torch |
|
from transformers import pipeline |
|
|
|
from PIL import Image |
|
|
|
import matplotlib.pyplot as plt |
|
|
|
from random import choice |
|
import io |
|
|
|
detector50 = pipeline(model="TuningAI/DETR-BASE_Marine") |
|
|
|
import gradio as gr |
|
|
|
COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff", |
|
"#7f7fff", "#7fbfff", "#7fffff", "#7fffbf", |
|
"#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"] |
|
|
|
fdic = { |
|
"family" : "Impact", |
|
"style" : "italic", |
|
"size" : 10, |
|
"color" : "red", |
|
"weight" : "bold" |
|
} |
|
|
|
|
|
def get_figure(in_pil_img, in_results): |
|
plt.figure(figsize=(16, 10)) |
|
plt.imshow(in_pil_img) |
|
ax = plt.gca() |
|
|
|
for prediction in in_results: |
|
selected_color = choice(COLORS) |
|
|
|
x, y = prediction['box']['xmin'], prediction['box']['ymin'], |
|
w, h = prediction['box']['xmax'] - prediction['box']['xmin'], prediction['box']['ymax'] - prediction['box']['ymin'] |
|
|
|
ax.add_patch(plt.Rectangle((x, y), w, h, fill=False, color=selected_color, linewidth=3)) |
|
ax.text(x, y, f"{prediction['label']}: {round(prediction['score']*100, 1)}%", fontdict=fdic) |
|
|
|
plt.axis("off") |
|
|
|
return plt.gcf() |
|
|
|
|
|
def infer(in_pil_img): |
|
results = detector50(in_pil_img) |
|
|
|
figure = get_figure(in_pil_img, results) |
|
|
|
buf = io.BytesIO() |
|
figure.savefig(buf, bbox_inches='tight') |
|
buf.seek(0) |
|
output_pil_img = Image.open(buf) |
|
|
|
return output_pil_img |
|
|
|
|
|
with gr.Blocks(title="DETR Object Detection") as demo: |
|
with gr.Row(): |
|
input_image = gr.Image(label="Input image", type="pil") |
|
output_image = gr.Image(label="Output image with predicted instances", type="pil") |
|
send_btn = gr.Button("Infer") |
|
send_btn.click(fn=infer, inputs=input_image, outputs=[output_image]) |
|
demo.launch(debug=True) |