YOLOv2 in PyTorch

This is a PyTorch implementation of YOLOv2.

This project is mainly based on darkflow and darknet.

For details about YOLO and YOLOv2 please refer to their project page and the paper:

YOLO9000: Better, Faster, Stronger by Joseph Redmon and Ali Farhadi.

I used a Cython extension for postprocessing and multiprocessing.Pool for image preprocessing.

Testing an image in VOC2007 costs about 13~20ms.

NOTE:

This is still an experimental project.

VOC07 test mAP is about 0.71 (trained on VOC07+12 trainval, reported by @cory8249).

See https://github.com/longcw/yolo2-pytorch/issues/1 and https://github.com/longcw/yolo2-pytorch/issues/23 for more details about training.

BTW, I recommend to write your own dataloader using torch.utils.data.Dataset since multiprocessing.Pool.imap won't stop even there is no enough memory space.

Installation and demo

1. Clone this repository

git clone git@github.com:longcw/yolo2-pytorch.git

2. Build the reorg layer (tf.extract_image_patches)

cd yolo2-pytorch

./make.sh

3. Download the trained model yolo-voc.weights.h5 and set the model path in demo.py

4. Run demo python demo.py.

Training YOLOv2

You can train YOLO2 on any dataset. Here we train it on VOC2007/2012.

1. Download the training, validation, test data and VOCdevkit

wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtrainval_06-Nov-2007.tar

wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCtest_06-Nov-2007.tar

wget http://host.robots.ox.ac.uk/pascal/VOC/voc2007/VOCdevkit_08-Jun-2007.tar

2. Extract all of these tars into one directory named VOCdevkit

tar xvf VOCtrainval_06-Nov-2007.tar

tar xvf VOCtest_06-Nov-2007.tar

tar xvf VOCdevkit_08-Jun-2007.tar

3. It should have this basic structure

$VOCdevkit/                          # development kit

$VOCdevkit/VOCcode/                  # VOC utility code

$VOCdevkit/VOC2007                    # image sets, annotations, etc.

# ... and several other directories ...

4. Since the program loading the data in yolo2-pytorch/data by default, you can set the data path as following.

cd yolo2-pytorch

mkdir data

cd data

ln -s $VOCdevkit VOCdevkit2007

5. Download the pretrained darknet19 model

and set the path in yolo2-pytorch/cfgs/exps/darknet19_exp1.py.

7. (optional) Training with TensorBoard.

To use the TensorBoard, install Crayon (https://github.com/torrvision/crayon)

and set use_tensorboard = True in yolo2-pytorch/cfgs/config.py.

6. Run the training program: python train.py.

Evaluation

Set the path of the trained_model in yolo2-pytorch/cfgs/config.py.

cd faster_rcnn_pytorch

mkdir output

python test.py