Training_Your_Model

You can get this Notebook on GitHub.

Step 1: Prepare you environment on your host computer

Note：This part of code run on your host computer

The Jupyter Notebook right up have a button like , then you choose Select Another Kernel, and choose Python Environments, then choose Creat Python Environment and choose Venv, then choose python3.10.

# Here is my hostcomputer information, you should install Ubantu 22.04 if you what use this code.
# Linux PC 6.8.0-45-generic #45~22.04.1-Ubuntu SMP PREEMPT_DYNAMIC Wed Sep 11 15:25:05 UTC 2 x86_64 x86_64 x86_64 GNU/Linux.

!uname -a

# Here is my python version, you should install python3.10.12 
# Python 3.10.12

!python -V

# Install libs, if you see 'install successfully' it means you install libs successfully, or when you see 'install error' it means you install libs unsuccessfully
try:
    %pip install ultralytics opencv-python matplotlib -q
    print('install successfully')
except Exception as e:
    print(f'install error: {e}')

Step 2: Prepare your dateset

I want to build a model to detect different fruit include banana, apple and orange. So I need to collect some pictures of this fruit. And I use robflow to label my dataset.

Step 1: Create Project

Select New Project:

Fill your project information and create the project

Step 2: Update images and annotate

Update images

Label your image

Step 3: Export dataset

Add annotated image to your dataset

Generate New version of your dataset

Download your dataset

# Download gdown to install dataset from google driver, if you see 'install successfully' it means you install libs successfully, or when you see 'install error' it means you install libs unsuccessfully
try:
    %pip install gdown -q
    print('install successfully')
except Exception as e:
    print(f'install error: {e}')

# Download your dataset, and you can also train your model on roboflow 
try:  
    !gdown https://drive.google.com/uc?id=1zZKnIVAcdNLUKg7IxaF-xLzE3Fvr3A05  && unzip roboflow.zip -d ~/datasets/ && rm roboflow.zip && mv ~/datasets/data.yaml ./data.yaml && cp -r ~/datasets/test/images ./
    print('install successfully')
except Exception as e:
    print(f'install error: {e}')

Step 3: Training Yolov8n

YOLOv8 is one of the most famous version in the YOLO (You Only Look Once) series developed by Ultralytics. It retains the key features of earlier versions, focusing on real-time object detection with improvements in speed, accuracy, and versatility across various tasks, such as object detection, segmentation, classification, and pose estimation.

# ! pip install ultralytics
! yolo settings

! yolo detect train data=data.yaml model=yolov8n.yaml epochs=100 imgsz=640

%matplotlib inline

import cv2
import os
import matplotlib.pyplot as plt
from ultralytics import YOLO


# Define the image directory path
image_dir = './images/'

# Load the best model
model = YOLO('./runs/detect/train/weights/best.pt')

# Get a list of all image files in the directory
image_files = [f for f in os.listdir(image_dir) if f.endswith(('.png', '.jpg', '.jpeg'))]

# Loop through each image in the directory
for image_file in image_files:
    image_path = os.path.join(image_dir, image_file)
    image = cv2.imread(image_path)

    # Perform prediction
    results = model(image)

    # Process the results and draw bounding boxes
    for result in results:
        # Extract bounding boxes
        boxes = result.boxes
        if boxes is not None:
            for box in boxes:
                conf = box.conf[0]  # Confidence score
                if conf > 0.5:
                    x1, y1, x2, y2 = map(int, box.xyxy[0])  # Get bounding box coordinates
                    cls = int(box.cls[0])  # Class label index
                    label = f"{model.names[cls]}: {conf:.2f}"  # Create label text

                    # Draw the bounding box and label on the image
                    cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)  # Green box
                    cv2.putText(image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)  # Green label

    # Convert BGR image (from OpenCV) to RGB for displaying with matplotlib
    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

    # Display the result image using matplotlib
    plt.figure(figsize=(10, 6))
    plt.imshow(image_rgb)
    plt.axis('off')  # Hide axis
    plt.title(f"YOLO Prediction - {image_file}")
    plt.show()

Step 4: Export to ONNX format

I will use Hailo DataFlow Compiler to convert model to hef format to inference on AI Kit, so I need to convert model to onnx format.

# Export Yolov11n model to onnx format, the path will be shown as below
# For me the path is /home/jiahao/Tutorial-of-AI-Kit-with-Raspberry-Pi-From-Zero-to-Hero/runs/detect/train/weights/best.onnx
model = YOLO('./runs/detect/train/weights/best.pt')
model.export(format="onnx", opset=9)