RayanRen/Object_detection_202

1import matplotlib.pyplot as plt2from PIL import Image3from transformers import DetrImageProcessor, DetrForObjectDetection4import torch5 6# colors for visualization7COLORS = [[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]]8 9import io10 11processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")12model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")13 14def fig2img(fig):15    buf = io.BytesIO()16    fig.savefig(buf)17    buf.seek(0)18    img = Image.open(buf)19    return img20 21def plot_results(image, results):22    plt.figure(figsize=(16, 10))23    plt.imshow(image)24    ax = plt.gca()25    colors = COLORS * 10026    for box, label, prob, color in zip(results["boxes"], results["labels"], results["scores"], colors):27        xmin, xmax, ymin, ymax = box[0].item(), box[2].item(), box[1].item(), box[3].item()28        ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,29                                   fill=False, color=color, linewidth=3))30        text = f'{model.config.id2label[label.item()]}: {prob:0.2f}'31        ax.text(xmin, ymin, text, fontsize=15,32                bbox=dict(facecolor='yellow', alpha=0.5))33    ax.axis("off")34    return fig2img(plt.gcf())35 36def predict(input_img):37    inputs = processor(images=input_img, return_tensors="pt")38    outputs = model(**inputs)39 40    target_sizes = torch.tensor([input_img.size[::-1]])41    results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.9)[0]42    return plot_results(input_img, results)43 44import gradio as gr45 46demo = gr.Interface(fn=predict,47                    inputs=gr.Image(type="pil"),48                    outputs="image")49demo.launch()