DEV Community

Cover image for Custom Object Detection on Jetson Orin Nano
Sampreeth Nataraj
Sampreeth Nataraj

Posted on

Custom Object Detection on Jetson Orin Nano

#nvidia #ai #machinelearning #ultralytics

How to train the model to identify custom objects?

This is the initial pre-trained model that has a set of objects and images for detection (it fails to recognize a pen)

This is after training it to recognize a pen

Here is a step by step guide on how to identify custom objects

I have used my NVIDIA Jetson Orin Nano running Jetpack7.2.1 for this project

  1. Verify your JetPack version. 
cat /etc/nv_tegra_release
Enter fullscreen mode Exit fullscreen mode
  1. Make sure the following JetPack contents are installed.
  • CUDA
  • cuDNN
  • TensorRT
  • VPI
  • multimedia libraries
  • NVIDIA container runtime
  1. If you are not sure the JetPack is installed then type the following command: 
sudo apt update
sudo apt install nvidia-jetpack
Enter fullscreen mode Exit fullscreen mode
  1. Then run 
sudo reboot
Enter fullscreen mode Exit fullscreen mode
  1. Verify the installed CUDA version: 
nvcc --version
Enter fullscreen mode Exit fullscreen mode
  1. Check if you have Python installed 
python3 --version
Enter fullscreen mode Exit fullscreen mode
  1. Create a YOLO project 
mkdir -p ~/yolo-jetson
cd ~/yolo-jetson
Enter fullscreen mode Exit fullscreen mode
  1. Create a Python virtual environment 
python3 -m venv yolo-env
Enter fullscreen mode Exit fullscreen mode
  1. Activate it 
source yolo-env/bin/activate
Enter fullscreen mode Exit fullscreen mode
  1. Every time you open a new terminal you can activate it with : 
cd ~/yolo-jetson
source yolo-env/bin/activate
Enter fullscreen mode Exit fullscreen mode
  1. Upgrade PIP 
python3 -m pip install --upgrade pip
Enter fullscreen mode Exit fullscreen mode
  1. Install Ultralytics 
pip install ultralytics
Enter fullscreen mode Exit fullscreen mode
  1. Install the correct version of PyTorch for JetPack 7.2.1. 
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu130
Enter fullscreen mode Exit fullscreen mode
  1. Verify PyTorch 
python3
Enter fullscreen mode Exit fullscreen mode
  1. Then, 
import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))
Enter fullscreen mode Exit fullscreen mode 
      or
Enter fullscreen mode Exit fullscreen mode 
python3 -c "import torch; print('PyTorch:', torch.__version__); print('CUDA available:', torch.cuda.is_available()); print('GPU:', torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'N/A')"
Enter fullscreen mode Exit fullscreen mode
  1. Install ONNX Runtime GPU 
pip install https://github.com/ultralytics/assets/releases/download/v0.0.0/onnxruntime_gpu-1.24.0-cp312-cp312-linux_aarch64.whl
Enter fullscreen mode Exit fullscreen mode
  1. Run your first Ultralytics program 
yolo predict model=yolo26n.pt source="https://github.com/ultralytics/assets/releases/download/v0.0.0/bus.jpg"
Enter fullscreen mode Exit fullscreen mode
  1. After this works you can run YOLO on your image.Create your image folder. 
mkdir -p images
Enter fullscreen mode Exit fullscreen mode

  1. You can use images that are already annotated or annotate them yourseves using roboflow.

  2. If you decide to annotate them yourself,

  • Create your project on Roboflow or any other similar platforms.
  • Annotate the object you want to train precisely.

  • Remember to annotate around 250-400 images for better accuracy.

  • Export the whole image set in YOLOv8 format

  1. From the directory containing your dataset: 
yolo detect train \
    model=yolo26n.pt \
    data=pen-dataset/data.yaml \
    epochs=100 \
    imgsz=640 \
    device=0
Enter fullscreen mode Exit fullscreen mode
  1. Run 
yolo predict \
    model=runs/detect/train/weights/best.pt \
    source=some_new_pen_image.jpg
Enter fullscreen mode Exit fullscreen mode
  1. Export to TensorRT 
yolo export \
    model=models/best.pt \
    format=engine \
    half=True
Enter fullscreen mode Exit fullscreen mode
  1. Run TensorRT 
yolo predict \
    model=models/best.engine \
    source=images/test_pen.jpg \
    device=0
Enter fullscreen mode Exit fullscreen mode
  1. Try it out!! 
yolo predict model=models/best.pt source=0 device=0 show=True
Enter fullscreen mode Exit fullscreen mode 
                       or
Enter fullscreen mode Exit fullscreen mode 
yolo predict model=models/best.engine source=0 device=0 show=True
Enter fullscreen mode Exit fullscreen mode
  1. If it is the first test then run the test by running the following command: 
cd ~/yolo-jetson
source yolo-env/bin/activate
yolo predict model=models/best.pt source=0 device=0 show=True
Enter fullscreen mode Exit fullscreen mode

Top comments (1)

Collapse
 
nataraj_p profile image
Nataraj P •

Good going