Modify demo.py

This commit is contained in:
Zhongdao 2020-01-29 23:56:24 +08:00
parent 04e7e79818
commit dee692baf0

View file

@ -12,6 +12,7 @@ from utils.io import read_results
from utils.log import logger from utils.log import logger
from utils.timer import Timer from utils.timer import Timer
from utils.evaluation import Evaluator from utils.evaluation import Evaluator
from utils.parse_config import parse_model_cfg
import utils.datasets as datasets import utils.datasets as datasets
import torch import torch
from track import eval_seq from track import eval_seq
@ -22,6 +23,9 @@ def track(opt):
result_root = opt.output_root if opt.output_root!='' else '.' result_root = opt.output_root if opt.output_root!='' else '.'
mkdir_if_missing(result_root) mkdir_if_missing(result_root)
cfg_dict = parse_model_cfg(opt.cfg)
opt.img_size = [int(cfg_dict[0]['width']), int(cfg_dict[0]['height'])]
# run tracking # run tracking
timer = Timer() timer = Timer()
accs = [] accs = []
@ -48,7 +52,6 @@ if __name__ == '__main__':
parser = argparse.ArgumentParser(prog='demo.py') parser = argparse.ArgumentParser(prog='demo.py')
parser.add_argument('--cfg', type=str, default='cfg/yolov3.cfg', help='cfg file path') 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('--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('--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('--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('--nms-thres', type=float, default=0.4, help='iou threshold for non-maximum suppression')