关于pytorch的yolov3的video文件
代码来自 eriklindernoren /yolov3,由于作者没有给出video代码,故在这里给出。
参考中3的代码感觉有多余部分(比如rgb和bgr的转换),去掉了一些多余的。增加了fps计算
参考:
- GitHub的pytorch版YOLOv3 链接:https://github.com/eriklindernoren/PyTorch-YOLOv3
- YOLOv3 论文作者的github链接:https://github.com/pjreddie/darknet
- 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()