File size: 2,545 Bytes
69af19f |
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 |
#!/usr/bin/env python
import pathlib
import gradio as gr
import numpy as np
import PIL.Image
import torch
from sahi.prediction import ObjectPrediction
from sahi.utils.cv import visualize_object_predictions
from transformers import AutoImageProcessor, DetaForObjectDetection
DESCRIPTION = '# DETA (Detection Transformers with Assignment)'
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
MODEL_ID = 'jozhang97/deta-swin-large'
image_processor = AutoImageProcessor.from_pretrained(MODEL_ID)
model = DetaForObjectDetection.from_pretrained(MODEL_ID)
model.to(device)
@torch.inference_mode()
def run(image_path: str, threshold: float) -> np.ndarray:
image = PIL.Image.open(image_path)
inputs = image_processor(images=image, return_tensors='pt').to(device)
outputs = model(**inputs)
target_sizes = torch.tensor([image.size[::-1]])
results = image_processor.post_process_object_detection(
outputs, threshold=threshold, target_sizes=target_sizes)[0]
boxes = results['boxes'].cpu().numpy()
scores = results['scores'].cpu().numpy()
cat_ids = results['labels'].cpu().numpy().tolist()
preds = []
for box, score, cat_id in zip(boxes, scores, cat_ids):
box = np.round(box).astype(int)
cat_label = model.config.id2label[cat_id]
pred = ObjectPrediction(bbox=box,
category_id=cat_id,
category_name=cat_label,
score=score)
preds.append(pred)
res = visualize_object_predictions(np.asarray(image), preds)['image']
return res
with gr.Blocks(css='style.css') as demo:
gr.Markdown(DESCRIPTION)
with gr.Row():
with gr.Column():
image = gr.Image(label='Input image', type='filepath')
threshold = gr.Slider(label='Score threshold',
minimum=0,
maximum=1,
value=0.1,
step=0.01)
run_button = gr.Button('Run')
result = gr.Image(label='Result', type='numpy')
with gr.Row():
paths = sorted(pathlib.Path('images').glob('*.jpg'))
gr.Examples(examples=[[path.as_posix(), 0.1] for path in paths],
inputs=[image, threshold],
outputs=result,
fn=run,
cache_examples=True)
run_button.click(fn=run, inputs=[image, threshold], outputs=result)
demo.queue().launch()
|