File size: 5,452 Bytes
cb670b6 e174663 cb670b6 195912f cb670b6 |
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 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 |
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
# Images
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
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. car, truck, traffic light", 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,
examples=examples,
title="Object Detection with YOLOS",
description=description,
).launch() |