YOLOv4 Pytorch implementation with all freebies and specials and 15+ more exclusive improvements. Easy to use!
!!! For the jupyter notebook please install pytorch-lightning version 0.7.6
This is implementation of YOLOv4 object detection neural network on pytorch. I'll try to implement all features of original paper.
You can use video_demo.py to take a look at the original weights realtime OD detection. (Have 9 fps on my GTX1060 laptop!!!)
You can train your own model with mosaic augmentation for training. Guides how to do this are written below. Borders of images on some datasets are even hard to find.
You can make inference, guide bellow.
#YOU CAN USE TORCH HUB
m = torch.hub.load("VCasecnikovs/Yet-Another-YOLOv4-Pytorch", "yolov4", pretrained=True)
import model
#If you change n_classes from the pretrained, there will be caught one error, don't panic it is ok
#FROM SAVED WEIGHTS
m = model.YOLOv4(n_classes=1, weights_path="weights/yolov4.pth")
#AUTOMATICALLY DOWNLOAD PRETRAINED
m = model.YOLOv4(n_classes=1, pretrained=True)
You can use torch hub or you can download weights using from this link: https://drive.google.com/open?id=12AaR4fvIQPZ468vhm0ZYZSLgWac2HBnq
import dataset
d = dataset.ListDataset("train.txt", img_dir='images', labels_dir='labels', img_extensions=['.JPG'], train=True)
path, img, bboxes = d[0]
!!! You can use SplitDataset.ipynb to create train.txt and valid.txt
"train.txt" is file which consists with filepaths to image (images\primula\DSC02542.JPG)
img_dir - Folder with images labels_dir - Folder with txt files for annotation img_extensions - extensions if images
If you set train=False -> uses letterboxes If you set train=True -> HSV augmentations and mosaic
dataset has collate_function
# collate func example
y1 = d[0]
y2 = d[1]
paths_b, xb, yb = d.collate_fn((y1, y2))
# yb has 6 columns
Is a tensor of size (B, 6), where B is amount of boxes in all batch images.
y_hat, loss = m(xb, yb)
!!! y_hat is already resized anchors to image size bboxes
y_hat, _ = m(img_batch) #_ is (0, 0, 0)
import utils
from PIL import Image
path, img, bboxes = d[0]
img_with_bboxes = utils.get_img_with_bboxes(img, bboxes[:, 2:]) #Returns numpy array
Image.fromarray(img_with_bboxes)
anchors, loss = m(xb.cuda(), yb.cuda())
confidence_threshold = 0.05
iou_threshold = 0.5
bboxes, labels = utils.get_bboxes_from_anchors(anchors, confidence_threshold, iou_threshold, coco_dict) #COCO dict is id->class dictionary (f.e. 0->person)
#For first img
arr = utils.get_img_with_bboxes(xb[0].cpu(), bboxes[0].cpu(), resize=False, labels=labels[0])
Image.fromarray(arr)
In case if you missed:
Paper Yolo v4: https://arxiv.org/abs/2004.10934
Original repo: https://github.com/AlexeyAB/darknet#how-to-train-to-detect-your-custom-objects
@article{yolov4,
title={YOLOv4: YOLOv4: Optimal Speed and Accuracy of Object Detection},
author={Alexey Bochkovskiy, Chien-Yao Wang, Hong-Yuan Mark Liao},
journal = {arXiv},
year={2020}
}