Simple TensorFlow Lite and Picamera2 example

tensorflow_example.py

#!/usr/bin/python3

import tflite_runtime.interpreter as tflite

import sys
import os
import argparse

import cv2
import numpy as np
from PIL import Image
from PIL import ImageDraw

#from null_preview import *
from qt_gl_preview import *
from picamera2 import *

parser = argparse.ArgumentParser()
parser.add_argument('--model', help='Path of the detection model.', required=True)
args = parser.parse_args()

# Configure Picamera2
picam2 = Picamera2()
#preview = NullPreview(picam2)
preview = QtGlPreview(picam2)
config = picam2.preview_configuration()
picam2.configure(config)
picam2.start()

# Capture an Image
np_array = picam2.capture_array()
img = Image.fromarray(np_array)
draw = ImageDraw.Draw(img, 'RGBA')
np_array = np_array[:,:,:3]

# Configure TensorFlow
interpreter = tflite.Interpreter(model_path=args.model, num_threads=4)
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
height = input_details[0]['shape'][1]
width = input_details[0]['shape'][2]
floating_model = False
if input_details[0]['dtype'] == np.float32:
    floating_model = True

initial_h, initial_w, channels = np_array.shape

image = cv2.resize(np_array, (width, height))
input_data = np.expand_dims(image, axis=0)
if floating_model:
   input_data = (np.float32(input_data) - 127.5) / 127.5

# Run TensorFlow
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()

detected_boxes = interpreter.get_tensor(output_details[0]['index'])
detected_classes = interpreter.get_tensor(output_details[1]['index'])
detected_scores = interpreter.get_tensor(output_details[2]['index'])
num_boxes = interpreter.get_tensor(output_details[3]['index'])

# Handle Results
for i in range(int(num_boxes)):
   top, left, bottom, right = detected_boxes[0][i]
   classId = int(detected_classes[0][i])
   score = detected_scores[0][i]

   if score > 0.5:
       xmin = left * initial_w
       ymin = bottom * initial_h
       xmax = right * initial_w
       ymax = top * initial_h
       for i in  range(5):
          rect_start = (xmin - i, ymin - i)
          rect_end = (xmax + i, ymax + i)
          draw.rectangle((rect_start, rect_end), outline = (0,128,128,20), fill = (0,128,128,20))
   img.save("out.png")