Picamera2 库 + TensorFlow Lite 实现实时对象检测

Raspberry Pi OS Bullseye 版本发布时弃用的 Picamera 以新面貌回归了，那就是 Picamera2 库。

下面介绍使用树莓派摄像头模块，用 Picamera2 库和 TensorFlow Lite 实现实时对象检测。

安装 Picamera2 库

因为目前 Picamera2 库还处于预览版，因此暂时只能通过 GitHub 编译安装。

$ sudo apt update

$ sudo apt install -y libboost-dev

$ sudo apt install -y libgnutls28-dev openssl libtiff5-dev

$ sudo apt install -y qtbase5-dev libqt5core5a libqt5gui5 libqt5widgets5

$ sudo apt install -y meson

$ sudo pip3 install pyyaml ply

$ sudo pip3 install --upgrade meson

$ sudo apt install -y libglib2.0-dev libgstreamer-plugins-base1.0-dev

$ git clone --branch picamera2 https://github.com/raspberrypi/libcamera.git

$ cd libcamera

$ meson build --buildtype=release -Dpipelines=raspberrypi -Dipas=raspberrypi -Dv4l2=true -Dgstreamer=enabled -Dtest=false -Dlc-compliance=disabled -Dcam=disabled -Dqcam=enabled -Ddocumentation=disabled -Dpycamera=enabled

$ ninja -C build

$ sudo ninja -C build install

$ cd ~

$ git clone https://github.com/tomba/kmsxx.git

$ cd kmsxx

$ git submodule update --init

$ sudo apt install -y libfmt-dev libdrm-dev

$ meson build

$ ninja -C build

$ cd ~

$ sudo pip3 install pyopengl

$ sudo apt install python3-pyqt5

$ git clone https://git@github.com:raspberrypi/picamera2.git

$ sudo pip3 install opencv-python==4.4.0.46

$ sudo apt install -y libatlas-base-dev

$ sudo pip3 install numpy --upgrade

$ cd ~

$ git clone https://github.com/RaspberryPiFoundation/python-v4l2.git

要跑起来还需要设置 PYTHONPATH 运行环境。比如你需要将下面的内容添加到 .bashrc 文件中。

export PYTHONPATH=/home/pi/picamera2:/home/pi/libcamera/build/src/py:/home/pi/kmsxx/build/py:/home/pi/python-v4l2

安装 TensorFlow Lite

由于我们将通过 Python 代码进行推理而不是训练，因此我们可以安装轻量级 TensorFlow Lite 运行时库以及我们需要的其他一些东西：

$ sudo apt install build-essentials

$ sudo apt install git

$ sudo apt install libatlas-base-dev

$ sudo apt install python3-pip

$ pip3 install tflite-runtime

$ pip3 install opencv-python==4.4.0.46

$ pip3 install pillow

$ pip3 install numpy

开始使用 TensorFlow Lite

安装完所有东西之后，下面我们构建一个演示程序：寻找图像中的苹果和香蕉。

代码会启用摄像头，并将采集到的图像不断传给 TensorFlow 的图像缓冲区。TensorFlow 随之在图像上进行对象检测。如果检测到任何对象，将用矩形框进行标注。

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 ImageFont, ImageDraw



from qt_gl_preview import *

from picamera2 import *



normalSize = (640, 480)

lowresSize = (320, 240)



rectangles = []



def ReadLabelFile(file_path):

  with open(file_path, 'r') as f:

    lines = f.readlines()

  ret = {}

  for line in lines:

    pair = line.strip().split(maxsplit=1)

    ret[int(pair[0])] = pair[1].strip()

  return ret



def DrawRectangles(request):

   stream = request.picam2.stream_map["main"]

   fb = request.request.buffers[stream]

   with fb.mmap(0) as b:

       im = np.array(b, copy=False, dtype=np.uint8).reshape((normalSize[1],normalSize[0], 4))



       for rect in rectangles:

          print(rect)

          rect_start = (int(rect[0]*2) - 5, int(rect[1]*2) - 5)

          rect_end = (int(rect[2]*2) + 5, int(rect[3]*2) + 5)

          cv2.rectangle(im, rect_start, rect_end, (0,255,0,0))

          if len(rect) == 5:

            text = rect[4]

            font = cv2.FONT_HERSHEY_SIMPLEX

            cv2.putText(im, text, (int(rect[0]*2) + 10, int(rect[1]*2) + 10), font, 1, (255,255,255),2,cv2.LINE_AA)

       del im



def InferenceTensorFlow( image, model, output, label=None):

   global rectangles



   if label:

       labels = ReadLabelFile(label)

   else:

       labels = None



   interpreter = tflite.Interpreter(model_path=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



   rgb = cv2.cvtColor(image,cv2.COLOR_GRAY2RGB)

   initial_h, initial_w, channels = rgb.shape



   picture = cv2.resize(rgb, (width, height))



   input_data = np.expand_dims(picture, axis=0)

   if floating_model:

      input_data = (np.float32(input_data) - 127.5) / 127.5



   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'])



   rectangles = []

   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

          box = [xmin, ymin, xmax, ymax]

          rectangles.append(box)

          if labels:

              print(labels[classId], 'score = ', score)

              rectangles[-1].append(labels[classId])

          else:

              print ('score = ', score)



def main():

    parser = argparse.ArgumentParser()

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

    parser.add_argument('--label', help='Path of the labels file.')

    parser.add_argument('--output', help='File path of the output image.')

    args = parser.parse_args()



    if ( args.output):

      output_file = args.output

    else:

      output_file = 'out.jpg'



    if ( args.label ):

      label_file = args.label

    else:

      label_file = None



    picam2 = Picamera2()

    preview = QtGlPreview(picam2)

    config = picam2.preview_configuration(main={"size": normalSize},

                                          lores={"size": lowresSize, "format": "YUV420"})

    picam2.configure(config)



    stride = picam2.stream_configuration("lores")["stride"]

    picam2.request_callback = DrawRectangles



    picam2.start()



    while True:

        buffer = picam2.capture_buffer("lores")

        grey = buffer[:stride*lowresSize[1]].reshape((lowresSize[1], stride))

        result = InferenceTensorFlow( grey, args.model, output_file, label_file )



if __name__ == '__main__':

  main()

via