TensorFlow Lite Integration

Relevant source files

The following files were used as context for generating this wiki page: - [examples/tensorflow/real_time.py](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/real_time.py) - [examples/tensorflow/segmentation.py](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/segmentation.py) - [examples/tensorflow/yolo_v5_real_time_with_labels.py](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/yolo_v5_real_time_with_labels.py) - [examples/tensorflow/compositing.py](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/compositing.py) - [examples/tensorflow/real_time_with_labels.py](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/real_time_with_labels.py) - [examples/tensorflow/remove_background.py](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/remove_background.py)

The picamera2 library provides robust support for real-time computer vision applications by integrating with TensorFlow Lite (TFLite). This allows developers to run lightweight machine learning models directly on a Raspberry Pi, processing video streams from the camera for tasks such as object detection, image segmentation, and classification. The integration leverages picamera2's ability to provide multiple, simultaneous video streams, enabling an efficient pipeline where a low-resolution stream is used for ML inference and a high-resolution stream is used for display.

The examples provided in the examples/tensorflow/ directory demonstrate various use cases. These scripts serve as practical templates for building custom applications. They typically involve configuring the camera, running a TFLite model on each frame from a low-resolution stream, and overlaying the results onto a high-resolution preview window.

System Architecture

The core architecture for TFLite integration involves a dual-stream camera configuration. A low-resolution stream (lores) is fed into the TFLite interpreter for processing, while a high-resolution stream (main) is used for the user-facing preview. This separation is key to achieving real-time performance, as it avoids the overhead of downscaling high-resolution images in software before inference.

The following diagram illustrates the typical data flow:

This diagram shows the parallel streams from Picamera2. The low-resolution stream goes through the inference pipeline, and the results are used to draw visualizations on the high-resolution preview window.

Sources: examples/tensorflow/real_time.py, examples/tensorflow/segmentation.py

Camera Configuration for Inference

To optimize for real-time inference, a specific camera configuration is used. The create_preview_configuration method is called to set up two streams:

1. main: A higher-resolution stream for display (e.g., 640x480).
2. lores: A low-resolution stream whose dimensions match the input tensor requirements of the TFLite model (e.g., 320x240). This stream is typically in YUV420 format for efficiency.

# examples/tensorflow/real_time.py:120-122
config = picam2.create_preview_configuration(main={"size": normalSize},
                                             lores={"size": lowresSize, "format": "YUV420"})
picam2.configure(config)

The stride of the lores stream is retrieved after configuration to correctly interpret the buffer later.

Parameter	Description	Example Value	Source
main	Configuration dictionary for the main, high-resolution display stream.	{"size": (640, 480)}	[examples/tensorflow/real_time.py:121](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/real_time.py#L121)
lores	Configuration dictionary for the low-resolution inference stream.	{"size": (320, 240), "format": "YUV420"}	[examples/tensorflow/real_time.py:122](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/real_time.py#L122)

Sources: examples/tensorflow/real_time.py:120-122, examples/tensorflow/segmentation.py:133-134, examples/tensorflow/yolo_v5_real_time_with_labels.py:108-110

Inference Pipeline

The main application logic resides in a while True loop that continuously captures frames, performs inference, and displays the results.

Sources: examples/tensorflow/real_time.py:128-131, examples/tensorflow/segmentation.py:140-144

1. Frame Capture and Pre-processing

Inside the loop, a buffer from the lores stream is captured using picam2.capture_buffer("lores"). This buffer is a 1D NumPy array. It is then reshaped into a 2D grayscale image based on the stream's stride and height.

# examples/tensorflow/real_time.py:129-130
buffer = picam2.capture_buffer("lores")
grey = buffer[:stride * lowresSize[1]].reshape((lowresSize[1], stride))

This grayscale image is then passed to the inference function.

Sources: examples/tensorflow/real_time.py:129-130, examples/tensorflow/segmentation.py:141-142

2. TFLite Model Execution

A dedicated function, often named InferenceTensorFlow, handles the core ML processing. This function is responsible for loading the model, preparing the input tensor, running inference, and returning the results.

The general flow within this function is as follows:

This flowchart shows the steps to execute a TensorFlow Lite model on a captured image.

Sources: examples/tensorflow/real_time.py:73-114, examples/tensorflow/segmentation.py:46-91

Key Steps:

1. Initialize Interpreter: An instance of tflite.Interpreter is created with the provided model path.
   Source: [examples/tensorflow/real_time.py:79](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/real_time.py#L79)
2. Prepare Image: The input grayscale image is converted to RGB and resized to match the model's expected input dimensions.
   Source: [examples/tensorflow/real_time.py:89-93](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/real_time.py#L89-L93)
3. Normalize Input: If the model is a floating-point model, the input data is normalized (e.g., scaled to a [-1, 1] or [0, 1] range).
   Source: [examples/tensorflow/segmentation.py:68-69](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/segmentation.py#L68-L69)
4. Run Inference: The prepared image data is set as the input tensor, interpreter.invoke() is called, and the output tensor is retrieved.
   Source: [examples/tensorflow/real_time.py:100-102](https://github.com/przkos/picamera2/blob/HEAD/examples/tensorflow/real_time.py#L100-L102)

3. Post-processing and Visualization

After inference, the raw output from the model must be processed to be useful. This varies depending on the model's task.

Object Detection

For object detection models (like MobileNet SSD or YOLO), the output tensors contain information about bounding boxes, class IDs, and confidence scores. The code iterates through the detections, filters them based on a score threshold (e.g., score > 0.5), and stores the coordinates and labels of valid objects.

The yolo_v5_real_time_with_labels.py example includes a specific YOLOdetect function to parse the unique output format of YOLOv5 models.

Sources: examples/tensorflow/real_time.py:104-114, examples/tensorflow/yolo_v5_real_time_with_labels.py:78-88

Image Segmentation

For segmentation models (like DeepLabV3), the output is a tensor where each value corresponds to a class index for that pixel. The code finds the class with the highest probability for each pixel (np.argmax) and generates colored masks for each detected class.

Sources: examples/tensorflow/segmentation.py:84-91, examples/tensorflow/remove_background.py:48-55

Visualization Techniques

Two primary methods are used to display the results on the preview window:

1. Post-Callback Drawing (post_callback): A function (e.g., DrawRectangles) is registered as a post_callback on the Picamera2 object. This function is called for every frame and receives the request object. It uses MappedArray to get direct, writable access to the main stream's buffer and draws shapes (like rectangles and text) on it using OpenCV (cv2.rectangle, cv2.putText).

   # examples/tensorflow/real_time_with_labels.py:50-52
   def DrawRectangles(request):
       with MappedArray(request, "main") as m:
           for rect in rectangles:
               # ... cv2.rectangle drawing logic ...
   
   # examples/tensorflow/real_time_with_labels.py:125
   picam2.post_callback = DrawRectangles
   

   Sources: examples/tensorflow/real_time_with_labels.py:50-61, 125

2. Overlay (set_overlay): For segmentation masks or other full-frame effects, an RGBA NumPy array is created and passed to picam2.set_overlay(). This overlay is rendered on top of the preview by the display system. This is useful for effects like background removal or displaying semi-transparent masks.

   # examples/tensorflow/segmentation.py:145-149
   overlay = np.zeros((normalSize[1], normalSize[0], 4), dtype=np.uint8)
   for v in masks.values():
       overlay += v
   overlay[:, :, -1][overlay[:, :, -1] == 255] = 150
   picam2.set_overlay(overlay)
   

   Sources: examples/tensorflow/segmentation.py:145-149, examples/tensorflow/remove_background.py:90-93

Interactive Compositing

The compositing.py example demonstrates a more advanced, interactive workflow. Instead of continuous inference, it waits for user input (Enter key) to capture a frame and run object detection. It then generates a binary mask for each detected object and saves both the original image and the mask to disk.

After capturing several objects, the user can specify which captured objects to composite together onto a new background, using PIL.Image.composite. This showcases how picamera2 can be used in interactive applications that combine ML with image manipulation.

This sequence diagram illustrates the interactive capture and compositing flow.

Sources: examples/tensorflow/compositing.py:134-141, 153-167

Summary of Key Functions and Arguments

The TensorFlow examples share a common structure, relying on a few key functions and command-line arguments.

Common Arguments

Argument	Description	Required	Example File
--model	Path to the .tflite model file.	Yes	real_time.py
--label	Path to the text file containing class labels.	No	real_time_with_labels.py
--colours	Path to a text file with colors for segmentation masks.	No	segmentation.py
--output	Path for the output image file (in relevant examples).	No	compositing.py

Sources: examples/tensorflow/real_time.py:117-119, examples/tensorflow/segmentation.py:124-128

Core Functions

Function	Description	Example File
InferenceTensorFlow	Takes a NumPy array, runs it through the TFLite interpreter, and processes the output.	real_time.py
ReadLabelFile	Reads a label map file into a dictionary, mapping class IDs to names.	real_time_with_labels.py
DrawRectangles	A post_callback function that draws bounding boxes on the main stream.	compositing.py
YOLOdetect	A specialized function to parse the output from a YOLOv5 model.	yolo_v5_real_time_with_labels.py

Sources: examples/tensorflow/real_time.py:73, examples/tensorflow/real_time_with_labels.py:41, examples/tensorflow/compositing.py:38, examples/tensorflow/yolo_v5_real_time_with_labels.py:78