Path: blob/master/site/en-snapshot/hub/tutorials/tf2_object_detection.ipynb
25118 views
Copyright 2020 The TensorFlow Hub Authors.
Licensed under the Apache License, Version 2.0 (the "License");
View on TensorFlow.org
Run in Google Colab
View on GitHub
Download notebook
See TF Hub modelsTensorFlow Hub Object Detection Colab
Welcome to the TensorFlow Hub Object Detection Colab! This notebook will take you through the steps of running an "out-of-the-box" object detection model on images.
Imports and Setup
Let's start with the base imports.
Utilities
Run the following cell to create some utils that will be needed later:
Helper method to load an image
Map of Model Name to TF Hub handle
List of tuples with Human Keypoints for the COCO 2017 dataset. This is needed for models with keypoints.
Visualization tools
To visualize the images with the proper detected boxes, keypoints and segmentation, we will use the TensorFlow Object Detection API. To install it we will clone the repo.
Intalling the Object Detection API
Now we can import the dependencies we will need later
Load label map data (for plotting).
Label maps correspond index numbers to category names, so that when our convolution network predicts 5, we know that this corresponds to airplane. Here we use internal utility functions, but anything that returns a dictionary mapping integers to appropriate string labels would be fine.
We are going, for simplicity, to load from the repository that we loaded the Object Detection API code
Build a detection model and load pre-trained model weights
Here we will choose which Object Detection model we will use. Select the architecture and it will be loaded automatically. If you want to change the model to try other architectures later, just change the next cell and execute following ones.
Tip: if you want to read more details about the selected model, you can follow the link (model handle) and read additional documentation on TF Hub. After you select a model, we will print the handle to make it easier.
Loading the selected model from TensorFlow Hub
Here we just need the model handle that was selected and use the Tensorflow Hub library to load it to memory.
Loading an image
Let's try the model on a simple image. To help with this, we provide a list of test images.
Here are some simple things to try out if you are curious:
Try running inference on your own images, just upload them to colab and load the same way it's done in the cell below.
Modify some of the input images and see if detection still works. Some simple things to try out here include flipping the image horizontally, or converting to grayscale (note that we still expect the input image to have 3 channels).
Be careful: when using images with an alpha channel, the model expect 3 channels images and the alpha will count as a 4th.
Doing the inference
To do the inference we just need to call our TF Hub loaded model.
Things you can try:
Print out
result['detection_boxes']and try to match the box locations to the boxes in the image. Notice that coordinates are given in normalized form (i.e., in the interval [0, 1]).inspect other output keys present in the result. A full documentation can be seen on the models documentation page (pointing your browser to the model handle printed earlier)
Visualizing the results
Here is where we will need the TensorFlow Object Detection API to show the squares from the inference step (and the keypoints when available).
the full documentation of this method can be seen here
Here you can, for example, set min_score_thresh to other values (between 0 and 1) to allow more detections in or to filter out more detections.
[Optional]
Among the available object detection models there's Mask R-CNN and the output of this model allows instance segmentation.
To visualize it we will use the same method we did before but adding an aditional parameter: instance_masks=output_dict.get('detection_masks_reframed', None)