02_yolo_object_detection

Chess, rolls or basketball? Let's create a custom object detection model

My posts on the Medium are usually very theoretical - I tend to analyse and describe the algorithms that define how Neural Networks work. This time, however, I decided to break this trend and show my readers how easy it is to train your own YOLO model, capable of detecting any objects we choose. In order to achieve this goal, we will need help from a very useful and easy-to-use implementation of YOLO. In short, not much coding, but a huge effect.

Setup environment

LINUX users can set up the environment with a single command. The script will download the YOLO implementation for you, and install all dependencies inside the Python virtual environment.

# Run in 02_yolo_object_detection directory
bash setup_yolo_framework.sh

Windows and MacOS users should use Docker to build their environment.

# Run in 02_yolo_object_detection directory
docker build -t yolo -f Dockerfile .
# Run docker image in interactive mode
docker run -it yolo:latest

Detection using pre-trained chess model

# Run in 02_yolo_object_detection directory
bash download_chess_model.sh
cd yolov3
# Run with activated virtual environment
python3 detect.py \
  --cfg cfg/chess.cfg \
  --weights ../chess/weights/best.pt \
  --source ../chess/samples \
  --names data/chess.names

Figure 1. Chess detection using TinyYOLO

Detection using pre-trained rolls model

# Run in 02_yolo_object_detection directory
bash download_rolls_model.sh
cd yolov3
# Run with activated virtual environment
python3 detect.py \
  --cfg cfg/rolls.cfg \
  --weights ../rolls/weights/best.pt \
  --source ../rolls/samples \
  --names data/rolls.names

Figure 2. Roll's detection using TinyYOLO

Detection using pre-trained basketball model

# Run in 02_yolo_object_detection directory
bash download_basketball_model.sh
cd yolov3
# Run with activated virtual environment
python3 detect.py \
  --cfg cfg/basketball.cfg \
  --weights ../basketball/weights/best.pt \
  --source ../basketball/samples \
  --names data/basketball.names

Figure 3. Detection of players moving around the basketball court,
based on YouTube-8M dataset.

Download chess dataset and train model on your own

# Run in 02_yolo_object_detection directory
bash download_chess_dataset.sh
cd yolov3
# Run with activated virtual environment
python3 train.py \
  --data data/chess.data \
  --cfg cfg/chess.cfg \
  --weights weights/yolo3.pt \
  --epochs 100