I’m trying to run a YOLOv8 segmentation model on a Raspberry Pi 5 equipped with an Edge TPU. I trained the model using yolov8n-seg on custom data.

During fine-tuning, I used the following command:

!yolo task=segment mode=train model=yolov8n-seg.pt data=/content/datasets/butterfly_data/butterfly_data.yaml epochs=100 imgsz=240 plots=True patience=0

butterfly_data.yaml:

train: /content/datasets/butterfly_data/train/
val: /content/datasets/butterfly_data/val/
nc: 1
names: ['butterfly']

I successfully converted the resulting best.pt file to best_full_integer_quant_edgetpu.tflite using the following command:

!yolo export model="/content/runs/segment/train/weights/best.pt" format="edgetpu" task=segment data=/content/datasets/butterfly_data/butterfly_data.yaml

The conversion went smoothly, and when I tested the best.pt model on my system without conversion, it worked perfectly. However, when I try to run the converted .tflite model on my Raspberry Pi 5 with the following code, I encounter some errors

test.py:

import cv2
import numpy as np
import pandas as pd
import cvzone
from ultralytics import YOLO
import time

# Load YOLOv8 segmentation model
# When It was best.pt, this code works perfectly
model = YOLO("best_full_integer_quant_edgetpu.tflite")

# Open webcam
cap = cv2.VideoCapture(0)

my_file = open("butterfly.txt", "r")
data = my_file.read()
class_list = data.split("n")

frame_count = 0
start_time = time.time()

while True:
    ret, frame = cap.read()
    if not ret:
        break
    
    frame_count += 1
    if frame_count % 3 != 0:
        continue

    # Resize frame
    frame = cv2.resize(frame, (1020, 600))
    overlay = frame.copy()

    results = model.predict(frame, imgsz=240)
    a = results[0].boxes.data
    px = pd.DataFrame(a).astype("float")
    for index, row in px.iterrows():
        x1 = int(row[0])
        y1 = int(row[1])
        x2 = int(row[2])
        y2 = int(row[3])
        
        d = int(row[5])
        c = class_list[d]
        cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
        cvzone.putTextRect(frame, f'{c}', (x1, y1), 1, 1)
    
    # Overlay segmentation mask on the original frame
    for result in results:
        if result.masks is not None:
           for j, mask in enumerate(result.masks.data):
               mask = mask.numpy() * 255
               mask = cv2.resize(mask.astype(np.uint8), (frame.shape[1], frame.shape[0]))
               mask_bgr = cv2.cvtColor(mask, cv2.COLOR_GRAY2BGR)
               contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
               cv2.fillPoly(mask_bgr, contours, (0, 0, 255))
               frame = cv2.addWeighted(frame, 1, mask_bgr, 0.5, 0)
    
    # Display frame
    end_time = time.time()
    elapsed_time = end_time - start_time
    fps = frame_count / elapsed_time
    cvzone.putTextRect(frame, f'FPS: {round(fps,2)}', (10, 30), 1, 1)

    cv2.imshow("YOLOv8 Segmentation", frame)
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

# Release video capture and close window
cap.release()
cv2.destroyAllWindows()

butterfly.txt :

butterfly

ERROR Message:

Traceback (most recent call last):
  File "/home/team00/yolov8_coral/segmenttflite.py", line 35, in <module>
    results = model.predict(frame, imgsz=256)
  File "/home/team00/yolov8_coral/.venv/lib/python3.9/site-packages/ultralytics/engine/model.py", line 567, in predict
    return self.predictor.predict_cli(source=source) if is_cli else self.predictor(source=source, stream=stream)
  File "/home/team00/yolov8_coral/.venv/lib/python3.9/site-packages/ultralytics/engine/predictor.py", line 168, in __call__
    return list(self.stream_inference(source, model, *args, **kwargs))  # merge list of Result into one
  File "/home/team00/yolov8_coral/.venv/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 35, in generator_context
    response = gen.send(None)
  File "/home/team00/yolov8_coral/.venv/lib/python3.9/site-packages/ultralytics/engine/predictor.py", line 274, in stream_inference
    s[i] += self.write_results(i, Path(paths[i]), im, s)
  File "/home/team00/yolov8_coral/.venv/lib/python3.9/site-packages/ultralytics/engine/predictor.py", line 334, in write_results
    string += f"{result.verbose()}{result.speed['inference']:.1f}ms"
  File "/home/team00/yolov8_coral/.venv/lib/python3.9/site-packages/ultralytics/engine/results.py", line 647, in verbose
    log_string += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, "
KeyError: 5

Interestingly, other models using object detection with Edge TPU work perfectly on my Raspberry Pi 5, but this segmentation model does not.

Before exporting, I tried installing different versions of TensorFlow (2.13.0 / 2.15.0) and retried the process.
The training and conversion of the model were all performed on Google Colab.