This repository contains code for a object detector based on YOLO9000: Better, Faster, Stronger, implementedin PyTorch. The code is based on the official code of YOLO v2, as well as a PyTorch port of the original code, by marvis. One of the goals of this code is to improve upon the original port by removing redundant parts of the code (The official code is basically a fully blown deep learning library, and includes stuff like sequence models, which are not used in YOLO). I've also tried to keep the code minimal, and document it as well as I can.
As of now, the code only contains the detection module, but you should expect the training module soon. :)
- Python 3.5
- OpenCV
- PyTorch 0.3+
Clone, and cd into the repo directory. Then, you have two variants of the detector, one that has been trained on PASCAL VOC
data (faster, but less accurate and recognises only 20 object categories), or the one trained on COCO (Slower, more accurate,
detects 80 categories).
For example, running the detector trained on PASCAL VOC download here, and place the weights file into your repo directory. Or, you could just type (if you're on Linux)
wget https://pjreddie.com/media/files/yolo-voc.weights
python detect.py --images imgs --det det --dataset pascal
For running with one trainined on coco, download this weightsfile, and run
the code with --dataset flag set to coco.
--images flag defines the directory to load images from, or a single image file (it will figure it out), and --det is the directory
to save images to. Other setting such as batch size, object threshold confidence can be tweaked with flags that can be looked up with
python detect.py -h
For this, you should run the file, video_demo.py with --video flag specifying the video file. The video file should be in .avi format since openCV only accepts OpenCV as the input format.
python video_demo.py --video video.avi --dataset pascal
Tweakable settings can be seen with -h flag.
To speed video inference, you can try using the video_demo_half.py file instead which does all the inference with 16-bit half precision floats instead of 32-bit float. I haven't seen big improvements, but I attribute that to having an older card (Tesla K80, Kepler arch, getting around 22 fps with PASCAL). If you have one of cards with fast float16 support, try it out, and if possible, benchmark it.
Same as video module, but you don't have to specify the video file since feed will be taken from your camera. To be precise, feed will be taken from what the OpenCV, recognises as camera 0. You can tweak this setting in code.
You'll have to download Tiny-yolo weightsfile in your repo folder. Excpect higher FPS and lower accuracy.
python cam_demo.py
You can easily tweak the code to use different weightsfiles, available at yolo website
Training module should arrive soon.