This project uses YOLOv8 for custom object detection. It's set up to train on custom datasets and perform predictions on images and videos.
NOTE: To protect the privacy of the data, the trained model weights as well as the training data is intentionally excluded from this repository.
- Clone this repository:
git clone https://github.com/bpk9/ai-object-detection.git
cd ai-object-detection- Install dependencies:
pip install -r requirements.txtai-object-detection/
├── datasets/ # Training and validation datasets
│ └── brian/ # Custom dataset (EX: "brian")
│ ├── train/ # Training images
│ ├── valid/ # Validation images
│ ├── test/ # Test images
│ └── data.yaml # Dataset configuration
├── train.py # Training script
└── predict.py # Prediction script
This project uses Roboflow for dataset preparation and annotation. To create your own dataset:
-
Create a free account on Roboflow
-
Create a new project:
- Select "Object Detection" as the project type
- Choose "YOLOv8" as the annotation format
-
Upload and annotate your images:
- Upload your images to Roboflow
- Use Roboflow's annotation tools to draw bounding boxes around objects
- Label each object with the appropriate class name
-
Generate and export your dataset:
- Add preprocessing steps if needed (resize, auto-orient, etc.)
- Add augmentations to increase dataset size (optional)
- Split your dataset (recommended: 70% train, 20% valid, 10% test)
- Export in YOLOv8 format
- Copy the exported dataset into your
datasets/directory - Update the path to your dataset in train.py
For more information on dataset preparation, visit Roboflow's YOLOv8 Guide.
The project uses YOLOv8 as the base model for training.
To train on your custom dataset:
-
Organize your dataset in the following structure:
datasets/<your-dataset-name>/train/images/: Training images (60%)datasets/<your-dataset-name>/valid/images/: Validation images (20%)datasets/<your-dataset-name>/test/images/: Test images (20%)
-
Update the
data_pathvariable in train.py -
Run the training script:
python train.pyTraining results will be saved in runs/detect/train/.
For additional documentation around training, see Model Training with Ultralytics YOLO
To run predictions on new images or videos:
-
Specify the path to your trained model weights in predict.py
-
Run the prediction script:
python predict.pyPrediction results will be saved in runs/detect/predict/.
For additional documentation around generating predictions, see Model Prediction with Ultralytics YOLO
The project is configured to use MPS (Metal Performance Shaders) for hardware acceleration on Apple Silicon Macs. Modify the device parameter in train.py if using different hardware:
- For NVIDIA GPU:
device="cuda" - For CPU:
device="cpu" - For Apple Silicon:
device="mps"