Skip to content

Open In Colab

IceVision Deployment App: COCO Dataset

This example uses Faster RCNN trained weights using the COCO dataset

About IceVision:

  • an Object-Detection Framework that connects to different libraries/frameworks such as Fastai, Pytorch Lightning, and Pytorch with more to come.

  • Features a Unified Data API with out-of-the-box support for common annotation formats (COCO, VOC, etc.)

  • Provides flexible model implementations with pluggable backbones

Installing packages

!pip install icevision[inference]

!pip install icedata

!pip install gradio

Imports

from icevision.all import *
import icedata
import PIL, requests
import torch
from torchvision import transforms
import gradio as gr

Loading trained model

class_map = icedata.coco.class_map()
model = icedata.coco.trained_models.faster_rcnn_resnet50_fpn()

Defininig the predict() method

def predict(
    model, image, detection_threshold: float = 0.5, mask_threshold: float = 0.5
):
    tfms_ = tfms.A.Adapter([tfms.A.Normalize()])
    # Whenever you have images in memory (numpy arrays) you can use `Dataset.from_images`
    infer_ds = Dataset.from_images([image], tfms_)

    batch, samples = faster_rcnn.build_infer_batch(infer_ds)
    preds = faster_rcnn.predict(
        model=model,
        batch=batch,
        detection_threshold=detection_threshold
    )
    return samples[0]["img"], preds[0]

Defining the show_preds method: called by gr.Interface(fn=show_preds, ...)

def show_preds(input_image, display_list, detection_threshold):
    display_label = ("Label" in display_list)
    display_bbox = ("BBox" in display_list)

    if detection_threshold==0: detection_threshold=0.5

    img, pred = predict(model=model, image=input_image, detection_threshold=detection_threshold)
    # print(pred)
    img = draw_pred(img=img, pred=pred, class_map=class_map, denormalize_fn=denormalize_imagenet, display_label=display_label, display_bbox=display_bbox)
    img = PIL.Image.fromarray(img)
    # print("Output Image: ", img.size, type(img))
    return img

Gradio User Interface

display_chkbox = gr.inputs.CheckboxGroup(["Label", "BBox"], label="Display")
detection_threshold_slider = gr.inputs.Slider(minimum=0, maximum=1, step=0.1, default=0.5, label="Detection Threshold")

outputs = gr.outputs.Image(type="pil")

gr_interface = gr.Interface(fn=show_preds, inputs=["image", display_chkbox,  detection_threshold_slider], outputs=outputs, title='IceApp - COCO')
gr_interface.launch(inline=False, share=True, debug=True)

Colab notebook detected. This cell will run indefinitely so that you can see errors and logs. To turn off, set debug=False in launch().
This share link will expire in 6 hours. If you need a permanent link, email support@gradio.app
Running on External URL: https://39710.gradio.app

Enjoy!

If you have any questions, please feel free to join us