truongginjs/atri-api repository overview

⁠YOLO Training & API Deployment Pipeline

This project provides a Docker-based pipeline for training a YOLO model using labeled image datasets and deploying an API for inference.

📌 GitHub Repository: truongginjs/atri-api⁠

⁠Table of Contents

• YOLO Training & API Deployment Pipeline⁠
  • Table of Contents⁠
  • Project Overview⁠
  • Prerequisites⁠
  • Dataset Preparation⁠
    • Directory Structure⁠
    • Label Format⁠
    • Class Labels⁠
    • Validation Checklist⁠
  • Running the Pipeline⁠
    • Sample shell script⁠
    • Basic Usage⁠
    • With GPU Support⁠
    • Skipping Specific Steps⁠
    • Running on Apple Silicon (ARM64)⁠
  • Environment Variables⁠
    • Optional Flags⁠
  • Script Behavior⁠
  • Container Management⁠
  • Troubleshooting⁠
    • 1. ONNX Runtime Error⁠
    • 2. CUDA Not Detected⁠
    • 3. YOLO Model Not Found⁠
  • Next Steps⁠
  • Contributors⁠

⁠Project Overview

This pipeline automates:

1. Data Preprocessing: Organizing images and labels in YOLO format.
2. Training: Training a YOLO model using Docker with optional GPU acceleration.
3. Model Export: Converting trained models into ONNX format for optimized inference.
4. Deployment: Running a containerized API server for real-time model inference.

📌 GitHub Repository: truongginjs/atri-api⁠

⁠Prerequisites

Ensure you have the following installed:

• Docker (for containerized execution)
• NVIDIA Drivers & CUDA (if using GPU acceleration)
• At least 8GB RAM (recommended)

⁠Dataset Preparation

⁠Directory Structure

Create an input/images/ folder with the following structure:

input/images/
├── image1.jpg
├── image1.txt
├── image2.jpg
├── image2.txt
└── ...

⁠Label Format

Labels follow the YOLO format:

<label_id> <x_center> <y_center> <width> <height>

Where:

• label_id: Integer class identifier
• x_center, y_center: Center coordinates (normalized 0-1)
• width, height: Bounding box dimensions (normalized 0-1)

Example:

0 0.656228 0.789873 0.163667 0.154430
1 0.440047 0.443038 0.610033 0.481013

⁠Class Labels

Define class labels in input/labels.yml:

labels:
  0: weight
  1: barcode

⁠Validation Checklist

✅ Ensure image and label files share the same name (e.g., image1.jpg ↔ image1.txt).
✅ Check that label values are normalized between 0 and 1.
✅ Each label file must contain at least one bounding box.
✅ Class IDs must match defined labels.

⁠Running the Pipeline

⁠Sample shell script

# full-pipeline.sh
#! /bin/bash

set -e

CURRENT_DIR=$(pwd)
TRAIN_IMAGE_NAME=truongginjs/atri-train
TRAIN_IMAGE_TAG=2.0-win

API_IMAGE_NAME=truongginjs/atri-api
API_IMAGE_TAG=2.0-win

# Add GPU flag parsing
USE_GPU=false
while getopts "g" opt; do
  case $opt in
    g) USE_GPU=true ;;
    *) echo "Usage: $0 [-g] (use -g for GPU support)" && exit 1 ;;
  esac
done

# Fix if statement syntax
if [ -d "${CURRENT_DIR}/output/yolo_train/train" ]; then
  echo "folder train found, remove that folder..."
  rm -rf ${CURRENT_DIR}/output/yolo_train/train
fi

echo "Starting training and exporting..."

# Define docker run command based on GPU flag
    # -e IGNORE_AUG=TRUE \
    # -e IGNORE_TRAIN=TRUE \
    # -e IGNORE_EXPORT=TRUE \
if [ "$USE_GPU" = true ]; then
  echo "Running with GPU support..."
  docker run --rm -it --gpus all --shm-size=16g \
    -v "${CURRENT_DIR}/input:/app/input" \
    -v "${CURRENT_DIR}/output:/app/output:rw" \
    -e EXPORT_PRE_TRAIN=/app/output/yolo_train/train/weights/best.pt \
    -e NUM_OF_AUGMENT=100 \
    -e EPOCHS=150 \
    -e IMGSZ=640 \
    ${TRAIN_IMAGE_NAME}:${TRAIN_IMAGE_TAG}
else
  echo "Running without GPU support..."
  docker run --rm -it \
    -v "${CURRENT_DIR}/input:/app/input" \
    -v "${CURRENT_DIR}/output:/app/output:rw" \
    -e EXPORT_PRE_TRAIN=/app/output/yolo_train/train/weights/best.pt \
    -e NUM_OF_AUGMENT=100 \
    -e EPOCHS=150 \
    -e IMGSZ=640 \
    ${TRAIN_IMAGE_NAME}:${TRAIN_IMAGE_TAG}
fi

mv ${CURRENT_DIR}/output/yolo_train/train/weights/best.onnx ${CURRENT_DIR}/output/best.onnx

echo "start api..."

# Check if container exists and remove it
CONTAINER_NAME="atri-api"
if [ "$(docker ps -q -f name=${CONTAINER_NAME})" ]; then
    echo "Stopping existing container..."
    docker stop ${CONTAINER_NAME}
fi
if [ "$(docker ps -aq -f name=${CONTAINER_NAME})" ]; then
    echo "Removing existing container..."
    docker rm ${CONTAINER_NAME}
fi

# Start new container
docker run -d -p 80:3000 \
  --name ${CONTAINER_NAME} \
  -v "${CURRENT_DIR}/output:/usr/src/app/models/yolo" \
  ${API_IMAGE_NAME}:${API_IMAGE_TAG}

⁠Basic Usage

To run the pipeline with CPU:

./full-pipeline.sh

⁠With GPU Support

If you have a compatible NVIDIA GPU:

./full-pipeline.sh -g

⁠Skipping Specific Steps

You can skip specific steps using environment variables:

• Skip data augmentation:

  ./full-pipeline.sh -e IGNORE_AUG=TRUE
  

  Copy
• Skip training:

  ./full-pipeline.sh -e IGNORE_TRAIN=TRUE
  

  Copy
• Skip multiple steps:

  ./full-pipeline.sh -e IGNORE_AUG=TRUE -e IGNORE_TRAIN=TRUE
  

  Copy
• Use GPU and skip augmentation:

  ./full-pipeline.sh -g -e IGNORE_AUG=TRUE
  

  Copy

⁠Running on Apple Silicon (ARM64)

Apple Silicon (M1/M2) does not support GPU acceleration for YOLO training inside Docker. Use:

./full-pipeline.arm64.sh

⁠Environment Variables

Modify these settings in the script:

⁠Optional Flags

docker run ... -e IGNORE_AUG=TRUE -e IGNORE_TRAIN=TRUE

⁠Script Behavior

1. Removes existing training data (if present).
2. Runs the training container with necessary volumes.
3. Exports the model to ONNX format.
4. Starts the API container on port 80.

⁠Container Management

The script automatically handles: ✅ Stopping and removing existing API containers.
✅ Starting a fresh container instance.

Manually stop the container:

docker stop atri-api

Check running containers:

docker ps

⁠Troubleshooting

⁠1. ONNX Runtime Error

If you see:

requirements: Ultralytics requirement ['onnxruntime-gpu'] not found, attempting AutoUpdate...

Manually install:

pip install onnxruntime-gpu

⁠2. CUDA Not Detected

Run:

nvidia-smi

If CUDA is missing, ensure NVIDIA drivers are installed.

⁠3. YOLO Model Not Found

Ensure that output/yolo_train/train/weights/best.pt exists before running the pipeline.

⁠Next Steps

✅ Train the model with new datasets
✅ Deploy the API for real-time inference
✅ Optimize training parameters for better accuracy

⁠Contributors

Developed by: TruongGinJS 🚀
Maintained by: TruongGinJS

📌 GitHub Repository: truongginjs/atri-api⁠