Skip to content

Latest commit

 

History

History
54 lines (45 loc) · 3.33 KB

README.md

File metadata and controls

54 lines (45 loc) · 3.33 KB
Raw

TensorFlow Lite Object Detection in Python

This code snipset is heavily based on TensorFlow Lite Object Detection
The detection model can be downloaded from above link.
For the realtime implementation on Android look into the Android Object Detection Example
Follow the object detection.ipynb to get information about how to use the TFLite model in your Python environment.

Details

The ssd_mobilenet_v1_1_metadata_1.tflite file's input takes normalized 300x300x3 shape image. And the output is composed of 4 different outputs. The 1st output contains the bounding box locations, 2nd output contains the label number of the predicted class, 3rd output contains the probabilty of the image belongs to the class, 4th output contains number of detected objects(maximum 10). The specific labels of the classes are stored in the labelmap.txt file.
I found the labelmap.txt in the Android Object Detection Example repository in below directory.

TFLite_examples/lite/examples/object_detection/android/app/src/main/assets

For model inference, we need to load, resize, typecast the image.
The mobileNet model uses uint8 format so typecast numpy array to uint8.

Then if you follow the correct instruction provided by Google in load_and_run_a_model_in_python, you would get output in below shape

Now we need to process this output to use it for object detection

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from PIL import Image

label_names = [line.rstrip('\n') for line in open("labelmap.txt")]
label_names = np.array(label_names)
numDetectionsOutput = int(np.minimum(numDetections[0],10))

for i in range(numDetectionsOutput):
    # Create figure and axes
    fig, ax = plt.subplots()

    # Display the image
    ax.imshow(res_im)

    # Create a Rectangle patch
    inputSize = 300
    left = outputLocations[0][i][1] * inputSize
    top = outputLocations[0][i][0] * inputSize
    right = outputLocations[0][i][3] * inputSize
    bottom = outputLocations[0][i][2] * inputSize
    class_name = label_names[int(outputClasses[0][i])]
    print("Output class: "+class_name+" | Confidence: "+ str(outputScores[0][i]))
    rect = patches.Rectangle((left, bottom), right-left, top-bottom, linewidth=1, edgecolor='r', facecolor='none')

    # Add the patch to the Axes
    ax.add_patch(rect)

    plt.show()


I believe you can modify the rest of the code as you want by yourself.
Thank you!