add video demo
This commit is contained in:
Zhongdao
parent
acc0a09d0f
commit
e47a5ae246
|
|
@ -1,3 +1,4 @@
|
|||
results/
|
||||
weights/
|
||||
data/
|
||||
tmp/
|
||||
|
|
|
|||
|
|
@ -0,0 +1,64 @@
|
|||
import os
|
||||
import os.path as osp
|
||||
import cv2
|
||||
import logging
|
||||
import argparse
|
||||
import motmetrics as mm
|
||||
|
||||
from tracker.multitracker import JDETracker
|
||||
from utils import visualization as vis
|
||||
from utils.utils import *
|
||||
from utils.log import logger
|
||||
from utils.timer import Timer
|
||||
from utils.evaluation import Evaluator
|
||||
import utils.datasets as datasets
|
||||
import torch
|
||||
from track import eval_seq
|
||||
|
||||
|
||||
def track(opt):
|
||||
logger.setLevel(logging.INFO)
|
||||
result_root = opt.output_root if opt.output_root!='' else '.'
|
||||
mkdir_if_missing(result_root)
|
||||
|
||||
# run tracking
|
||||
timer = Timer()
|
||||
accs = []
|
||||
n_frame = 0
|
||||
|
||||
logger.info('start tracking...')
|
||||
dataloader = datasets.LoadVideo(opt.input_video, opt.img_size)
|
||||
result_filename = os.path.join(result_root, 'results.txt')
|
||||
frame_rate = dataloader.frame_rate
|
||||
|
||||
frame_dir = None if opt.output_format=='text' else osp.join(result_root, 'frame')
|
||||
try:
|
||||
eval_seq(opt, dataloader, 'mot', result_filename,
|
||||
save_dir=frame_dir, show_image=False, frame_rate=frame_rate)
|
||||
except Exception as e:
|
||||
logger.info(e)
|
||||
|
||||
if opt.output_format == 'video':
|
||||
output_video_path = osp.join(result_root, 'result.mp4')
|
||||
cmd_str = 'ffmpeg -f image2 -i {}/%05d.jpg -c:v copy {}'.format(osp.join(result_root, 'frame'), output_video_path)
|
||||
os.system(cmd_str)
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(prog='demo.py')
|
||||
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
|
||||
parser.add_argument('--weights', type=str, default='weights/latest.pt', help='path to weights file')
|
||||
parser.add_argument('--img-size', type=int, default=(1088, 608), help='size of each image dimension')
|
||||
parser.add_argument('--iou-thres', type=float, default=0.5, help='iou threshold required to qualify as detected')
|
||||
parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold')
|
||||
parser.add_argument('--nms-thres', type=float, default=0.4, help='iou threshold for non-maximum suppression')
|
||||
parser.add_argument('--min-box-area', type=float, default=200, help='filter out tiny boxes')
|
||||
parser.add_argument('--track-buffer', type=int, default=30, help='tracking buffer')
|
||||
parser.add_argument('--input-video', type=str, help='path to the input video')
|
||||
parser.add_argument('--output-format', type=str, default='video', help='expected output format, can be video, tracklet or text')
|
||||
parser.add_argument('--output-root', type=str, default='results', help='expected output root path')
|
||||
opt = parser.parse_args()
|
||||
print(opt, end='\n\n')
|
||||
|
||||
track(opt)
|
||||
|
||||
15
track.py
15
track.py
|
|
@ -12,11 +12,8 @@ from utils.timer import Timer
|
|||
from utils.evaluation import Evaluator
|
||||
import utils.datasets as datasets
|
||||
import torch
|
||||
from utils.utils import *
|
||||
|
||||
def mkdirs(path):
|
||||
if os.path.exists(path):
|
||||
return
|
||||
os.makedirs(path)
|
||||
|
||||
def write_results(filename, results, data_type):
|
||||
if data_type == 'mot':
|
||||
|
|
@ -41,9 +38,7 @@ def write_results(filename, results, data_type):
|
|||
|
||||
|
||||
def eval_seq(opt, dataloader, data_type, result_filename, save_dir=None, show_image=True, frame_rate=30):
|
||||
if save_dir is not None:
|
||||
mkdirs(save_dir)
|
||||
|
||||
mkdir_if_missing(save_dir)
|
||||
tracker = JDETracker(opt, frame_rate=frame_rate)
|
||||
timer = Timer()
|
||||
results = []
|
||||
|
|
@ -85,7 +80,7 @@ def main(opt, data_root='/data/MOT16/train', det_root=None, seqs=('MOT16-05',),
|
|||
save_images=False, save_videos=False, show_image=True):
|
||||
logger.setLevel(logging.INFO)
|
||||
result_root = os.path.join(data_root, '..', 'results', exp_name)
|
||||
mkdirs(result_root)
|
||||
mkdir_if_missing(result_root)
|
||||
data_type = 'mot'
|
||||
|
||||
# run tracking
|
||||
|
|
@ -131,8 +126,7 @@ def main(opt, data_root='/data/MOT16/train', det_root=None, seqs=('MOT16-05',),
|
|||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(prog='test.py')
|
||||
parser.add_argument('--batch-size', type=int, default=8, help='size of each image batch')
|
||||
parser = argparse.ArgumentParser(prog='track.py')
|
||||
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path')
|
||||
parser.add_argument('--weights', type=str, default='weights/latest.pt', help='path to weights file')
|
||||
parser.add_argument('--img-size', type=int, default=(1088, 608), help='size of each image dimension')
|
||||
|
|
@ -140,7 +134,6 @@ if __name__ == '__main__':
|
|||
parser.add_argument('--conf-thres', type=float, default=0.5, help='object confidence threshold')
|
||||
parser.add_argument('--nms-thres', type=float, default=0.4, help='iou threshold for non-maximum suppression')
|
||||
parser.add_argument('--min-box-area', type=float, default=200, help='filter out tiny boxes')
|
||||
parser.add_argument('--pixel-mean', type=float, default=[0,0,0], nargs='+', help='pixel mean')
|
||||
parser.add_argument('--track-buffer', type=int, default=30, help='tracking buffer')
|
||||
parser.add_argument('--test-mot16', action='store_true', help='tracking buffer')
|
||||
parser.add_argument('--save-images', action='store_true', help='save tracking results (image)')
|
||||
|
|
|
|||
|
|
@ -75,6 +75,54 @@ class LoadImages: # for inference
|
|||
return self.nF # number of files
|
||||
|
||||
|
||||
class LoadVideo: # for inference
|
||||
def __init__(self, path, img_size=(1088, 608)):
|
||||
self.cap = cv2.VideoCapture(path)
|
||||
self.frame_rate = int(round(self.cap.get(cv2.CAP_PROP_FPS)))
|
||||
self.vw = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
self.vh = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
self.vn = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
|
||||
self.width = img_size[0]
|
||||
self.height = img_size[1]
|
||||
self.count = 0
|
||||
|
||||
self.w, self.h = self.get_size(self.vw, self.vh, self.width, self.height)
|
||||
print('Lenth of the video: {:d} frames'.format(self.vn))
|
||||
|
||||
def get_size(self, vw, vh, dw, dh):
|
||||
wa, ha = float(dw) / vw, float(dh) / vh
|
||||
a = min(wa, ha)
|
||||
return int(vw *a), int(vh*a)
|
||||
|
||||
def __iter__(self):
|
||||
self.count = -1
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
self.count += 1
|
||||
if self.count == len(self):
|
||||
raise StopIteration
|
||||
# Read image
|
||||
res, img0 = self.cap.read() # BGR
|
||||
assert img0 is not None, 'Failed to load frame {:d}'.format(self.count)
|
||||
img0 = cv2.resize(img0, (self.w, self.h))
|
||||
|
||||
# Padded resize
|
||||
img, _, _, _ = letterbox(img0, height=self.height, width=self.width)
|
||||
|
||||
# Normalize RGB
|
||||
img = img[:, :, ::-1].transpose(2, 0, 1)
|
||||
img = np.ascontiguousarray(img, dtype=np.float32)
|
||||
img /= 255.0
|
||||
|
||||
# cv2.imwrite(img_path + '.letterbox.jpg', 255 * img.transpose((1, 2, 0))[:, :, ::-1]) # save letterbox image
|
||||
return self.count, img, img0
|
||||
|
||||
def __len__(self):
|
||||
return self.vn # number of files
|
||||
|
||||
|
||||
class LoadImagesAndLabels: # for training
|
||||
def __init__(self, path, img_size=(1088,608), augment=False, transforms=None):
|
||||
with open(path, 'r') as file:
|
||||
|
|
|
|||
|
|
@ -52,27 +52,6 @@ def model_info(model): # Plots a line-by-line description of a PyTorch model
|
|||
print('Model Summary: %g layers, %g parameters, %g gradients\n' % (i + 1, n_p, n_g))
|
||||
|
||||
|
||||
def coco_class_weights(): # frequency of each class in coco train2014
|
||||
weights = 1 / torch.FloatTensor(
|
||||
[187437, 4955, 30920, 6033, 3838, 4332, 3160, 7051, 7677, 9167, 1316, 1372, 833, 6757, 7355, 3302, 3776, 4671,
|
||||
6769, 5706, 3908, 903, 3686, 3596, 6200, 7920, 8779, 4505, 4272, 1862, 4698, 1962, 4403, 6659, 2402, 2689,
|
||||
4012, 4175, 3411, 17048, 5637, 14553, 3923, 5539, 4289, 10084, 7018, 4314, 3099, 4638, 4939, 5543, 2038, 4004,
|
||||
5053, 4578, 27292, 4113, 5931, 2905, 11174, 2873, 4036, 3415, 1517, 4122, 1980, 4464, 1190, 2302, 156, 3933,
|
||||
1877, 17630, 4337, 4624, 1075, 3468, 135, 1380])
|
||||
weights /= weights.sum()
|
||||
return weights
|
||||
|
||||
|
||||
def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper)
|
||||
# https://tech.amikelive.com/node-718/what-object-categories-labels-are-in-coco-dataset/
|
||||
# a = np.loadtxt('data/coco.names', dtype='str', delimiter='\n')
|
||||
# b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n')
|
||||
# x = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco
|
||||
x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34,
|
||||
35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
|
||||
64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90]
|
||||
return x
|
||||
|
||||
|
||||
def plot_one_box(x, img, color=None, label=None, line_thickness=None): # Plots one bounding box on image img
|
||||
tl = line_thickness or round(0.0004 * max(img.shape[0:2])) + 1 # line thickness
|
||||
|
|
@ -505,28 +484,6 @@ def strip_optimizer_from_checkpoint(filename='weights/best.pt'):
|
|||
torch.save(a, filename.replace('.pt', '_lite.pt'))
|
||||
|
||||
|
||||
def coco_class_count(path='../coco/labels/train2014/'):
|
||||
# histogram of occurrences per class
|
||||
|
||||
nC = 80 # number classes
|
||||
x = np.zeros(nC, dtype='int32')
|
||||
files = sorted(glob.glob('%s/*.*' % path))
|
||||
for i, file in enumerate(files):
|
||||
labels = np.loadtxt(file, dtype=np.float32).reshape(-1, 5)
|
||||
x += np.bincount(labels[:, 0].astype('int32'), minlength=nC)
|
||||
print(i, len(files))
|
||||
|
||||
|
||||
def coco_only_people(path='../coco/labels/val2014/'):
|
||||
# find images with only people
|
||||
|
||||
files = sorted(glob.glob('%s/*.*' % path))
|
||||
for i, file in enumerate(files):
|
||||
labels = np.loadtxt(file, dtype=np.float32).reshape(-1, 5)
|
||||
if all(labels[:, 0] == 0):
|
||||
print(labels.shape[0], file)
|
||||
|
||||
|
||||
def plot_results():
|
||||
# Plot YOLO training results file 'results.txt'
|
||||
# import os; os.system('wget https://storage.googleapis.com/ultralytics/yolov3/results_v1.txt')
|
||||
|
|
|
|||
Loading…
Reference in New Issue