|
from transformers import AutoFeatureExtractor, YolosForObjectDetection |
|
import gradio as gr |
|
from PIL import Image |
|
import torch |
|
import matplotlib.pyplot as plt |
|
import io |
|
import numpy as np |
|
import os |
|
os.system("pip -qq install yoloxdetect==0.0.7") |
|
from yoloxdetect import YoloxDetector |
|
|
|
|
|
torch.hub.download_url_to_file('https://tochkanews.ru/wp-content/uploads/2020/09/0.jpg', '1.jpg') |
|
torch.hub.download_url_to_file('https://s.rdrom.ru/1/pubs/4/35893/1906770.jpg', '2.jpg') |
|
torch.hub.download_url_to_file('https://static.mk.ru/upload/entities/2022/04/17/07/articles/detailPicture/5b/39/28/b6/ffb1aa636dd62c30e6ff670f84474f75.jpg', '3.jpg') |
|
|
|
def yolox_inference( |
|
image_path: gr.inputs.Image = None, |
|
model_path: gr.inputs.Dropdown = 'kadirnar/yolox_s-v0.1.1', |
|
config_path: gr.inputs.Textbox = 'configs.yolox_s', |
|
image_size: gr.inputs.Slider = 640 |
|
): |
|
""" |
|
YOLOX inference function |
|
Args: |
|
image: Input image |
|
model_path: Path to the model |
|
config_path: Path to the config file |
|
image_size: Image size |
|
Returns: |
|
Rendered image |
|
""" |
|
|
|
model = YoloxDetector(model_path, config_path=config_path, device="cpu", hf_model=True) |
|
pred = model.predict(image_path=image_path, image_size=image_size) |
|
return pred |
|
|
|
|
|
inputs = [ |
|
gr.inputs.Image(type="filepath", label="Input Image"), |
|
gr.inputs.Dropdown( |
|
label="Model Path", |
|
choices=[ |
|
"kadirnar/yolox_s-v0.1.1", |
|
"kadirnar/yolox_m-v0.1.1", |
|
"kadirnar/yolox_tiny-v0.1.1", |
|
], |
|
default="kadirnar/yolox_s-v0.1.1", |
|
), |
|
gr.inputs.Dropdown( |
|
label="Config Path", |
|
choices=[ |
|
"configs.yolox_s", |
|
"configs.yolox_m", |
|
"configs.yolox_tiny", |
|
], |
|
default="configs.yolox_s", |
|
), |
|
gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"), |
|
] |
|
|
|
outputs = gr.outputs.Image(type="filepath", label="Output Image") |
|
title = "YOLOX is a high-performance anchor-free YOLO." |
|
|
|
examples = [ |
|
["1.jpg", "kadirnar/yolox_m-v0.1.1", "configs.yolox_m", 640], |
|
["2.jpg", "kadirnar/yolox_s-v0.1.1", "configs.yolox_s", 640], |
|
["3.jpg", "kadirnar/yolox_tiny-v0.1.1", "configs.yolox_tiny", 640], |
|
] |
|
|
|
demo_app = gr.Interface( |
|
fn=yolox_inference, |
|
inputs=inputs, |
|
outputs=outputs, |
|
title=title, |
|
examples=examples, |
|
cache_examples=True, |
|
theme='huggingface', |
|
) |
|
demo_app.launch(debug=True, enable_queue=True) |
|
|
|
|
|
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 get_class_list_from_input(classes_string: str): |
|
if classes_string == "": |
|
return [] |
|
classes_list = classes_string.split(",") |
|
classes_list = [x.strip() for x in classes_list] |
|
return classes_list |
|
|
|
def infer(img, model_name: str, prob_threshold: int, classes_to_show = str): |
|
feature_extractor = AutoFeatureExtractor.from_pretrained(f"hustvl/{model_name}") |
|
model = YolosForObjectDetection.from_pretrained(f"hustvl/{model_name}") |
|
|
|
img = Image.fromarray(img) |
|
|
|
pixel_values = feature_extractor(img, return_tensors="pt").pixel_values |
|
|
|
with torch.no_grad(): |
|
outputs = model(pixel_values, output_attentions=True) |
|
|
|
probas = outputs.logits.softmax(-1)[0, :, :-1] |
|
keep = probas.max(-1).values > prob_threshold |
|
|
|
target_sizes = torch.tensor(img.size[::-1]).unsqueeze(0) |
|
postprocessed_outputs = feature_extractor.post_process(outputs, target_sizes) |
|
bboxes_scaled = postprocessed_outputs[0]['boxes'] |
|
|
|
classes_list = get_class_list_from_input(classes_to_show) |
|
res_img = plot_results(img, probas[keep], bboxes_scaled[keep], model, classes_list) |
|
|
|
return res_img |
|
|
|
def plot_results(pil_img, prob, boxes, model, classes_list): |
|
plt.figure(figsize=(16,10)) |
|
plt.imshow(pil_img) |
|
ax = plt.gca() |
|
colors = COLORS * 100 |
|
for p, (xmin, ymin, xmax, ymax), c in zip(prob, boxes.tolist(), colors): |
|
cl = p.argmax() |
|
object_class = model.config.id2label[cl.item()] |
|
|
|
if len(classes_list) > 0 : |
|
if object_class not in classes_list: |
|
continue |
|
|
|
ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, |
|
fill=False, color=c, linewidth=3)) |
|
text = f'{object_class}: {p[cl]:0.2f}' |
|
ax.text(xmin, ymin, text, fontsize=15, |
|
bbox=dict(facecolor='yellow', alpha=0.5)) |
|
plt.axis('off') |
|
return fig2img(plt.gcf()) |
|
|
|
def fig2img(fig): |
|
buf = io.BytesIO() |
|
fig.savefig(buf) |
|
buf.seek(0) |
|
img = Image.open(buf) |
|
return img |
|
|
|
description = """Object Detection with YOLOS. Choose https://github.com/amikelive/coco-labels/blob/master/coco-labels-2014_2017.txtyour model and you're good to go. |
|
You can adapt the minimum probability threshold with the slider. |
|
Additionally you can restrict the classes that will be shown by putting in a comma separated list of |
|
[COCO classes](https://github.com/amikelive/coco-labels/blob/master/coco-labels-2014_2017.txt). |
|
Leaving the field empty will show all classes""" |
|
|
|
image_in = gr.components.Image() |
|
image_out = gr.components.Image() |
|
model_choice = gr.components.Dropdown(["yolos-tiny", "yolos-small", "yolos-base", "yolos-small-300", "yolos-small-dwr"], value="yolos-small", label="YOLOS Model") |
|
prob_threshold_slider = gr.components.Slider(minimum=0, maximum=1.0, step=0.01, value=0.9, label="Probability Threshold") |
|
classes_to_show = gr.components.Textbox(placeholder="e.g. person, boat", label="Classes to use (empty means all classes)") |
|
|
|
Iface = gr.Interface( |
|
fn=infer, |
|
inputs=[image_in,model_choice, prob_threshold_slider, classes_to_show], |
|
outputs=image_out, |
|
|
|
title="Object Detection with YOLOS", |
|
description=description, |
|
).launch() |
