关于pytorch的yolov3的video文件

参考：

1. GitHub的pytorch版YOLOv3 链接：https://github.com/eriklindernoren/PyTorch-YOLOv3
2. YOLOv3 论文作者的github链接：https://github.com/pjreddie/darknet
3. Guardian-Li给出的video代码：https://github.com/eriklindernoren/PyTorch-YOLOv3/issues/485

附上代码

from __future__ import division

from models import *
from utils.utils import *
from utils.datasets import *

import time
import argparse
import cv2
import torch
from torch.autograd import Variable


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--image_folder", type=str, default="data/samples", help="path to dataset")
    parser.add_argument("--video_file", type=str, default=r"C:\Users\DELL\Desktop\project\opencv_combat\walking_people.mp4", help="path to dataset")
    parser.add_argument("--model_def", type=str, default="config/yolov3.cfg", help="path to model definition file")
    parser.add_argument("--weights_path", type=str, default="weights/yolov3.weights", help="path to weights file")
    parser.add_argument("--class_path", type=str, default="data/coco.names", help="path to class label file")
    parser.add_argument("--conf_thres", type=float, default=0.6, help="object confidence threshold")
    parser.add_argument("--nms_thres", type=float, default=0.4, help="iou thresshold for non-maximum suppression")
    parser.add_argument("--batch_size", type=int, default=1, help="size of the batches")
    parser.add_argument("--n_cpu", type=int, default=3, help="number of cpu threads to use during batch generation")
    parser.add_argument("--img_size", type=int, default=416, help="size of each image dimension")
    parser.add_argument("--checkpoint_model", type=str, help="path to checkpoint model")
    opt = parser.parse_args()
    print(opt)


    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = Darknet(opt.model_def, img_size=opt.img_size).to(device)
    if opt.weights_path.endswith(".weights"):
        # Load darknet weights
        model.load_darknet_weights(opt.weights_path)
    else:
        # Load checkpoint weights
        model.load_state_dict(torch.load(opt.weights_path))
    model.eval()  # Set in evaluation mode
    classes = load_classes(opt.class_path)
    Tensor = torch.cuda.FloatTensor if torch.cuda.is_available() else torch.FloatTensor


    #if opt.vedio_file.endswith(".mp4"):
    cap = cv2.VideoCapture(opt.video_file)
    colors = np.random.randint(0, 255, size=(len(classes), 3), dtype="uint8")

    start = time.time()
    frames = 0
    #NUM=0
    while cap.isOpened():
        ret, img = cap.read()
        if ret is False:
            break
        # img = cv2.resize(img, (1280, 960), interpolation=cv2.INTER_CUBIC)
        # RGBimg=changeBGR2RGB(img)
        imgTensor = transforms.ToTensor()(img)
        imgTensor, _ = pad_to_square(imgTensor, 0)
        imgTensor = resize(imgTensor, opt.img_size)
        #加一个维度
        imgTensor = imgTensor.unsqueeze(0)
        imgTensor = Variable(imgTensor.type(Tensor))
        with torch.no_grad():
            detections = model(imgTensor)
            detections = non_max_suppression(detections, opt.conf_thres, opt.nms_thres)

        if detections[0] is  None:
            pass
        else:
            # output is a list and have one element，we can achieve the array
            output = detections[0]
            detections = rescale_boxes(output, opt.img_size, img.shape[:2])
            unique_labels = detections[:, -1].cpu().unique()
            n_cls_preds = len(unique_labels)
            for x1, y1, x2, y2, conf, cls_conf, cls_pred in detections:
                box_w = x2 - x1
                box_h = y2 - y1
                color = [int(c) for c in colors[int(cls_pred)]]
                #print(cls_conf)
                img = cv2.rectangle(img, (x1, y1 + box_h), (x2, y1), color, 2)
                cv2.putText(img, classes[int(cls_pred)], (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
                cv2.putText(img, str("%.2f" % float(conf)), (x2, y2 - box_h), cv2.FONT_HERSHEY_SIMPLEX, 0.5,
                            color, 2)
        #计算fps
        frames += 1
        print("FPS of the video is {:5.4f}".format(frames / (time.time() - start)))
        cv2.imshow('frame',img)
        #cv2.waitKey(0)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()