File size: 2,268 Bytes
301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab 301fa4e 5b802ab | 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 | import gradio as gr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import torch
from transformers import AutoImageProcessor, AutoModelForObjectDetection
repo_id = "magomerob/yolo_finetuned_raccoons"
processor = AutoImageProcessor.from_pretrained(repo_id)
model = AutoModelForObjectDetection.from_pretrained(repo_id)
model.eval()
def show_preds(input_image, display_label=True, display_bbox=True, detection_threshold=0.5):
if detection_threshold == 0:
detection_threshold = 0.5
if isinstance(input_image, np.ndarray):
img = Image.fromarray(input_image).convert("RGB")
else:
img = input_image.convert("RGB")
inputs = processor(images=img, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
# convert model outputs to boxes in image space
target_sizes = torch.tensor([img.size[::-1]]) # (h, w)
results = processor.post_process_object_detection(
outputs, threshold=float(detection_threshold), target_sizes=target_sizes
)[0]
draw = ImageDraw.Draw(img)
try:
font = ImageFont.truetype("DejaVuSans.ttf", 16)
except Exception:
font = ImageFont.load_default()
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
x1, y1, x2, y2 = [int(round(v)) for v in box.tolist()]
name = model.config.id2label.get(int(label), str(int(label)))
s = float(score)
if display_bbox:
draw.rectangle([x1, y1, x2, y2], width=3)
if display_label:
text = f"{name} {s:.2f}"
tw, th = draw.textbbox((0, 0), text, font=font)[2:]
pad = 3
draw.rectangle([x1, max(0, y1 - th - 2*pad), x1 + tw + 2*pad, y1], fill="black")
draw.text((x1 + pad, max(0, y1 - th - pad)), text, font=font, fill="white")
return img
demo = gr.Interface(
fn=show_preds,
inputs=[
gr.Image(type="numpy"),
gr.Checkbox(label="Label", value=True),
gr.Checkbox(label="Box", value=True),
gr.Slider(0, 1, step=0.05, value=0.5, label="Detection Threshold"),
],
outputs=gr.Image(type="pil"),
examples=[["raccoon-101.jpg", True, True, 0.5]],
)
if __name__ == "__main__":
demo.launch() |