def main():
global args, logger, v_id
args = parser.parse_args()
cfg = load_config(args)
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info(args)
# setup model
if args.arch == 'Custom':
from custom import Custom
model = Custom(anchors=cfg['anchors'])
else:
parser.error('invalid architecture: {}'.format(args.arch))
if args.resume:
assert isfile(args.resume), '{} is not a valid file'.format(args.resume)
model = load_pretrain(model, args.resume)
model.eval()
device = torch.device('cuda' if (torch.cuda.is_available() and not args.cpu) else 'cpu')
model = model.to(device)
# setup dataset
dataset = load_dataset(args.dataset)
# VOS or VOT?
if args.dataset in ['DAVIS','DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask:
vos_enable = True # enable Mask output
else:
vos_enable = False
total_lost = 0 # VOT
iou_lists = [] # VOS
speed_list = []
for v_id, video in enumerate(dataset.keys(), start=1):
if args.video != '' and video != args.video:
continue
if vos_enable:
iou_list, speed = track_vos(model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask, args.refine, args.dataset in ['DAVIS2017', 'ytb_vos'], device=device)
iou_lists.append(iou_list)
else:
lost, speed = track_vot(model, dataset[video], cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask, args.refine, device=device)
total_lost += lost
speed_list.append(speed)
# report final result
if vos_enable:
for thr, iou in zip(thrs, np.mean(np.concatenate(iou_lists), axis=0)):
logger.info('Segmentation Threshold {:.2f} mIoU: {:.3f}'.format(thr, iou))
else:
logger.info('Total Lost: {:d}'.format(total_lost))
logger.info('Mean Speed: {:.2f} FPS'.format(np.mean(speed_list)))
def __init__(self, config_path, model_path):
args = TrackArgs()
args.config = config_path
args.resume = model_path
cfg = load_config(args)
if args.arch == 'Custom':
from custom import Custom
self.model = Custom(anchors=cfg['anchors'])
else:
parser.error('invalid architecture: {}'.format(args.arch))
if args.resume:
assert isfile(args.resume), '{} is not a valid file'.format(args.resume)
self.model = load_pretrain(self.model, args.resume)
self.model.eval()
self.device = torch.device('cuda' if (torch.cuda.is_available() and not args.cpu) else 'cpu')
self.model = self.model.to(self.device)
################# Dangerous
self.p = TrackerConfig()
self.p.update(cfg['hp'] if 'hp' in cfg.keys() else None, self.model.anchors)
self.p.renew()
self.p.scales = self.model.anchors['scales']
self.p.ratios = self.model.anchors['ratios']
self.p.anchor_num = self.model.anchor_num
self.p.anchor = generate_anchor(self.model.anchors, self.p.score_size)
if self.p.windowing == 'cosine':
self.window = np.outer(np.hanning(self.p.score_size), np.hanning(self.p.score_size))
elif self.p.windowing == 'uniform':
self.window = np.ones((self.p.score_size, self.p.score_size))
self.window = np.tile(self.window.flatten(), self.p.anchor_num)
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args() # args通过解析获得的
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global') # 实例化一个记录器
logger.info("\n" + collect_env_info())
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(
cfg, indent=4))) # 转变成json格式的文件，缩进4格
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True, anchors=cfg['anchors'])
else:
exit()
logger.info(model)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
model = model.cuda() # 模型转移到GPU上
dist_model = torch.nn.DataParallel(
model, list(range(torch.cuda.device_count()))).cuda() # 多GPU训练
if args.resume and args.start_epoch != 0: # 这是在干啥？蒙蔽了！！！！！
model.features.unfix((args.start_epoch - 1) / args.epochs)
optimizer, lr_scheduler = build_opt_lr(model, cfg, args,
args.start_epoch) # 如何构建优化器和学习策略？？？
# optionally resume from a checkpoint
if args.resume:
assert os.path.isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model, optimizer, args.start_epoch, best_acc, arch = restore_from(
model, optimizer, args.resume)
dist_model = torch.nn.DataParallel(
model, list(range(torch.cuda.device_count()))).cuda()
logger.info(lr_scheduler)
logger.info('model prepare done')
train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch,
cfg)
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
init_log('global', logging.INFO) # 返回一个logger对象,logging_INFO是日志的等级
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global') # 获取上面初始化的logger对象
logger.info("\n" + collect_env_info())
logger.info(args)
cfg = load_config(args) # 返回修改后的配置文件对象
logger.info("config \n{}".format(json.dumps(cfg, indent=4))) #json.loads()是将str转化成dict格式，json.dumps()是将dict转化成str格式。
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True, anchors=cfg['anchors'])
else:
exit()
logger.info(model)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
model = model.cuda()
dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()
if args.resume and args.start_epoch != 0:
model.features.unfix((args.start_epoch - 1) / args.epochs)
optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch)
# optionally resume from a checkpoint
if args.resume:
assert os.path.isfile(args.resume), '{} is not a valid file'.format(args.resume)
model, optimizer, args.start_epoch, best_acc, arch = restore_from(model, optimizer, args.resume)
dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()
logger.info(lr_scheduler)
logger.info('model prepare done')
train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg)
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
args = args_process(args)
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info("\n" + collect_env_info())
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
# build dataset
train_loader, val_loader = build_data_loader(cfg)
args.img_size = int(cfg['train_datasets']['search_size'])
args.nms_threshold = float(cfg['train_datasets']['RPN_NMS'])
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True,
opts=args,
anchors=train_loader.dataset.anchors)
else:
exit()
logger.info(model)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
else:
raise Exception("Pretrained weights must be loaded!")
model = model.cuda()
dist_model = torch.nn.DataParallel(model,
list(range(
torch.cuda.device_count()))).cuda()
logger.info('model prepare done')
logger = logging.getLogger('global')
val_avg = AverageMeter()
validation(val_loader, dist_model, cfg, val_avg)
def __init__(self, args):
super(PatchTrainer, self).__init__()
# Setup device
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.backends.cudnn.benchmark = True
# Setup tracker cfg
cfg = load_config(args)
p = TrackerConfig()
p.renew()
self.p = p
# Setup tracker
siammask = Tracker(p=p, anchors=cfg['anchors'])
if args.resume:
assert isfile(args.resume), 'Please download {} first.'.format(args.resume)
siammask = load_pretrain(siammask, args.resume)
siammask.eval().to(self.device)
self.model = siammask
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info("\n" + collect_env_info())
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
args = args_process(args)
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info("\n" + collect_env_info())
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
args.img_size = int(cfg['train_datasets']['search_size'])
args.nms_threshold = float(cfg['train_datasets']['RPN_NMS'])
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True,
opts=args,
anchors=train_loader.dataset.anchors)
else:
exit()
logger.info(model)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
model = model.cuda()
dist_model = torch.nn.DataParallel(model,
list(range(
torch.cuda.device_count()))).cuda()
if args.resume and args.start_epoch != 0:
model.features.unfix((args.start_epoch - 1) / args.epochs)
optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch)
# optionally resume from a checkpoint
if args.resume:
assert os.path.isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model, optimizer, args.start_epoch, best_acc, arch = restore_from(
model, optimizer, args.resume)
dist_model = torch.nn.DataParallel(
model, list(range(torch.cuda.device_count()))).cuda()
logger.info(lr_scheduler)
logger.info('model prepare done')
global cur_lr
if not os.path.exists(args.save_dir): # makedir/save model
os.makedirs(args.save_dir)
num_per_epoch = len(train_loader.dataset) // args.batch
num_per_epoch_val = len(val_loader.dataset) // args.batch
for epoch in range(args.start_epoch, args.epochs):
lr_scheduler.step(epoch)
cur_lr = lr_scheduler.get_cur_lr()
logger = logging.getLogger('global')
train_avg = AverageMeter()
val_avg = AverageMeter()
if dist_model.module.features.unfix(epoch / args.epochs):
logger.info('unfix part model.')
optimizer, lr_scheduler = build_opt_lr(dist_model.module, cfg,
args, epoch)
train(train_loader, dist_model, optimizer, lr_scheduler, epoch, cfg,
train_avg, num_per_epoch)
if dist_model.module.features.unfix(epoch / args.epochs):
logger.info('unfix part model.')
optimizer, lr_scheduler = build_opt_lr(dist_model.module, cfg,
args, epoch)
if (epoch + 1) % args.save_freq == 0:
save_checkpoint(
{
'epoch': epoch,
'arch': args.arch,
'state_dict': dist_model.module.state_dict(),
'best_acc': best_acc,
'optimizer': optimizer.state_dict(),
'anchor_cfg': cfg['anchors']
}, False,
os.path.join(args.save_dir, 'checkpoint_e%d.pth' % (epoch)),
os.path.join(args.save_dir, 'best.pth'))
validation(val_loader, dist_model, epoch, cfg, val_avg,
num_per_epoch_val)
default='../../data/tennis',
help='datasets')
parser.add_argument('--gt_file',
default=None,
type=str,
help='ground truth txt file')
parser.add_argument('--cpu', action='store_true', help='cpu mode')
args = parser.parse_args()
if __name__ == '__main__':
# Setup device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.backends.cudnn.benchmark = True
# Setup Model
cfg = load_config(args)
siammask = Custom_(anchors=cfg['anchors'])
if args.resume:
assert isfile(args.resume), 'Please download {} first.'.format(
args.resume)
siammask = load_pretrain(siammask, args.resume)
siammask.eval().to(device)
# Parse Image file
img_files = sorted(glob.glob(join(args.base_path, '*.jp*')))
ims = [cv2.imread(imf) for imf in img_files]
# Select ROI
cv2.namedWindow("SiamMask", cv2.WND_PROP_FULLSCREEN)
# cv2.setWindowProperty("SiamMask", cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
def main():
global args, logger, v_id
args = parser.parse_args()
cfg = load_config(args)
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info(args)
# setup model
if args.arch == 'Custom':
from custom import Custom
model = Custom(anchors=cfg['anchors'])
else:
parser.error('invalid architecture: {}'.format(args.arch))
if args.resume:
assert isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model = load_pretrain(model, args.resume)
model.eval()
device = torch.device('cuda' if (
torch.cuda.is_available() and not args.cpu) else 'cpu')
model = model.to(device)
# setup dataset
dataset = load_dataset(args.dataset, args.dir_type)
# VOS or VOT?
if args.dataset in ['DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask:
vos_enable = True # enable Mask output
else:
vos_enable = False
# total_lost = 0 # VOT
# iou_lists = [] # VOS
# speed_list = []
for v_id, video in enumerate(dataset.keys(), start=1):
if args.video != '' and video != args.video:
continue
if vos_enable:
iou_list, speed = track_vos(
model,
dataset[video],
cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask,
args.refine,
args.dataset in ['DAVIS2017', 'ytb_vos'],
device=device)
# iou_lists.append(iou_list)
else:
lost, speed = track_vot(model,
dataset[video],
cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask,
args.refine,
device=device)
total_lost += lost
def process_vedio(vedio_path, initRect):
"""
视频处理
:param vedio_path：视频路径
:param initRect: 跟踪目标的初始位置
:return:
"""
# 1. 设置设备信息 Setup device
# 有GPU时选择GPU，否则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 默认优化运行效率
torch.backends.cudnn.benchmark = True
# 2. 模型设置 Setup Model
# 2.1 将命令行参数解析出来
cfg = load_config(args)
# 2.2 custom是构建的网络，否则引用model中的网络结构
from custom import Custom
siammask = Custom(anchors=cfg['anchors'])
# 2.3 判断是否存在模型的权重文件
if args.resume:
assert isfile(args.resume), 'Please download {} first.'.format(
args.resume)
siammask = load_pretrain(siammask, args.resume)
# 在运行推断前，需要调用 model.eval() 函数，以将 dropout 层 和 batch normalization 层设置为评估模式(非训练模式).
# to(device)将张量复制到GPU上，之后的计算将在GPU上运行
siammask.eval().to(device)
# 首帧跟踪目标的位置
x, y, w, h = initRect
print(x)
VeryBig = 999999999 # 用于将视频框调整到最大
Cap = cv2.VideoCapture(vedio_path) # 设置读取摄像头
ret, frame = Cap.read() # 读取帧
ims = [frame] # 把frame放入列表格式的frame， 因为原文是将每帧图片放入列表
im = frame
f = 0
target_pos = np.array([x + w / 2, y + h / 2])
target_sz = np.array([w, h])
state = siamese_init(im, target_pos, target_sz, siammask,
cfg['hp']) # init tracker"
middlepath = "../data/middle.mp4"
outpath = "../data/output.mp4"
vediowriter = cv2.VideoWriter(middlepath,
cv2.VideoWriter_fourcc('M', 'P', '4', 'V'),
10, (320, 240))
while (True):
tic = cv2.getTickCount()
ret, im = Cap.read() # 逐个提取frame
if (ret == False):
break
state = siamese_track(state, im, mask_enable=True,
refine_enable=True) # track
location = state['ploygon'].flatten()
mask = state['mask'] > state['p'].seg_thr
im[:, :, 2] = (mask > 0) * 255 + (mask == 0) * im[:, :, 2]
cv2.polylines(im, [np.int0(location).reshape((-1, 1, 2))], True,
(0, 255, 0), 3)
vediowriter.write(im)
cv2.imshow('SiamMask', im)
key = cv2.waitKey(1)
if key > 0:
break
f = f + 1
vediowriter.release()
return
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
print("Init logger")
logger = logging.getLogger('global')
print(44)
#logger.info("\n" + collect_env_info())
print(99)
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
print(2)
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
print(3)
path = "/usr4/alg504/cliao25/siammask/experiments/siammask_base/snapshot/checkpoint_e{}.pth"
for epoch in range(1,21):
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True, anchors=cfg['anchors'])
else:
exit()
print(4)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
model = model.cuda()
#model.features.unfix((epoch - 1) / 20)
optimizer, lr_scheduler = build_opt_lr(model, cfg, args, epoch)
filepath = path.format(epoch)
assert os.path.isfile(filepath)
model, _, _, _, _ = restore_from(model, optimizer, filepath)
#model = load_pretrain(model, filepath)
model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()
model.train()
device = torch.device('cuda')
model = model.to(device)
valid(val_loader, model, cfg)
print("Done")
def main():
# 获取命令行参数信息
global args, logger, v_id
args = parser.parse_args()
# 获取配置文件中配置信息：主要包括网络结构，超参数等
cfg = load_config(args)
# 初始化logxi信息，并将日志信息输入到磁盘文件中
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
# 将相关的配置信息输入到日志文件中
logger = logging.getLogger('global')
logger.info(args)
# setup model
# 加载网络模型架构
if args.arch == 'Custom':
from custom import Custom
model = Custom(anchors=cfg['anchors'])
else:
parser.error('invalid architecture: {}'.format(args.arch))
# 加载网络模型参数
if args.resume:
assert isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model = load_pretrain(model, args.resume)
# 使用评估模式，将drop等激活
model.eval()
# 硬件信息
device = torch.device('cuda' if (
torch.cuda.is_available() and not args.cpu) else 'cpu')
model = model.to(device)
# 加载数据集 setup dataset
dataset = load_dataset(args.dataset)
# 这三种数据支持掩膜 VOS or VOT?
if args.dataset in ['DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask:
vos_enable = True # enable Mask output
else:
vos_enable = False
total_lost = 0 # VOT
iou_lists = [] # VOS
speed_list = []
# 对数据进行处理
for v_id, video in enumerate(dataset.keys(), start=1):
if args.video != '' and video != args.video:
continue
# true 调用track_vos
if vos_enable:
# 如测试数据是['DAVIS2017', 'ytb_vos']时，会开启多目标跟踪
iou_list, speed = track_vos(
model,
dataset[video],
cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask,
args.refine,
args.dataset in ['DAVIS2017', 'ytb_vos'],
device=device)
iou_lists.append(iou_list)
# False 调用track_vot
else:
lost, speed = track_vot(model,
dataset[video],
cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask,
args.refine,
device=device)
total_lost += lost
speed_list.append(speed)
# report final result
if vos_enable:
for thr, iou in zip(thrs, np.mean(np.concatenate(iou_lists), axis=0)):
logger.info('Segmentation Threshold {:.2f} mIoU: {:.3f}'.format(
thr, iou))
else:
logger.info('Total Lost: {:d}'.format(total_lost))
logger.info('Mean Speed: {:.2f} FPS'.format(np.mean(speed_list)))
def main():
# args.base_path = base_path
args.resume = "../SiamMask/experiments/siammask_sharp/SiamMask_DAVIS.pth"
args.config = "../SiamMask/experiments/siammask_sharp/config_davis.json"
print(join(args.base_path, 'groundtruth_rect.txt'))
# Setup device
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
torch.backends.cudnn.benchmark = True
# Setup Model
cfg = load_config(args)
p = TrackerConfig()
p.renew()
siammask = Tracker(p=p, anchors=cfg['anchors'])
if args.resume:
assert isfile(args.resume), 'Please download {} first.'.format(args.resume)
siammask = load_pretrain(siammask, args.resume)
siammask.eval().to(device)
# Parse Image file
img_files = sorted(glob.glob(join(join(args.base_path, 'imgs'), '*.jp*')))
ims = [cv2.imread(imf) for imf in img_files]
# Select ROI
cv2.namedWindow("SiamMask", cv2.WND_PROP_FULLSCREEN)
try:
init_rect = cv2.selectROI('SiamMask', ims[0], False, False)
gts = None
x, y, w, h = init_rect
except:
exit()
file1 = open(join(args.base_path, 'groundtruth_rect.txt'), 'w')
file1.write('{0:d},{1:d},{2:d},{3:d}\n'.format(x, y, w, h))
toc = 0
for f, im in enumerate(ims):
tic = cv2.getTickCount()
if f == 0: # init
target_pos = np.array([x + w / 2, y + h / 2])
target_sz = np.array([w, h])
state = tracker_init(im, target_pos, target_sz, siammask, device=device) # init tracker
state['gts'] = gts
state['device'] = device
elif f > 0: # tracking
state = tracker_track(state, im, siammask, device=device) # track
target_pos, target_sz =state['target_pos'], state['target_sz']
x, y = (target_pos - target_sz/2).astype(int)
x2, y2 = (target_pos + target_sz/2).astype(int)
cv2.rectangle(im, (x, y), (x2, y2), (0, 255, 0), 4)
cv2.imshow('SiamMask', im)
key = cv2.waitKey(1)
if key == ord('q'):
break
file1.write('{0:d},{1:d},{2:d},{3:d}\n'.format(x, y, x2-x, y2-y))
toc += cv2.getTickCount() - tic
file1.close()
toc /= cv2.getTickFrequency()
fps = f / toc
print('SiamMask Time: {:02.1f}s Speed: {:3.1f}fps (with visulization!)'.format(toc, fps))
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
args = args_process(args)
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info("\n" + collect_env_info())
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
args.img_size = int(cfg['train_datasets']['search_size'])
args.nms_threshold = float(cfg['train_datasets']['RPN_NMS'])
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True,
opts=args,
anchors=train_loader.dataset.anchors)
else:
exit()
logger.info(model)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
model = model.cuda()
dist_model = torch.nn.DataParallel(model,
list(range(
torch.cuda.device_count()))).cuda()
if args.resume and args.start_epoch != 0:
model.features.unfix((args.start_epoch - 1) / args.epochs)
optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch)
# optionally resume from a checkpoint
if args.resume:
assert os.path.isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model, optimizer, args.start_epoch, best_acc, arch = restore_from(
model, optimizer, args.resume)
dist_model = torch.nn.DataParallel(
model, list(range(torch.cuda.device_count()))).cuda()
logger.info(lr_scheduler)
logger.info('model prepare done')
train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch,
cfg)
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
logger.info("\n" + collect_env_info())
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True, anchors=cfg['anchors'])
else:
model = models.__dict__[args.arch](anchors=cfg['anchors'])
logger.info(model)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
model = model.cuda()
dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()
if args.resume and args.start_epoch != 0:
model.features.unfix((args.start_epoch - 1) / args.epochs)
optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch)
logger.info(lr_scheduler)
# optionally resume from a checkpoint
if args.resume:
assert os.path.isfile(args.resume), '{} is not a valid file'.format(args.resume)
model, optimizer, args.start_epoch, best_acc, arch = restore_from(model, optimizer, args.resume)
dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()
epoch = args.start_epoch
if dist_model.module.features.unfix(epoch/args.epochs):
logger.info('unfix part model.')
optimizer, lr_scheduler = build_opt_lr(dist_model.module, cfg, args, epoch)
lr_scheduler.step(epoch)
cur_lr = lr_scheduler.get_cur_lr()
logger.info('epoch:{} resume lr {}'.format(epoch, cur_lr))
logger.info('model prepare done')
train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch, cfg)
def main():
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
params = {'penalty_k': args.penalty_k,
'window_influence': args.window_influence,
'lr': args.lr,
'instance_size': args.search_region}
num_search = len(params['penalty_k']) * len(params['window_influence']) * \
len(params['lr']) * len(params['instance_size'])
print(params)
print(num_search)
cfg = load_config(args)
if args.arch == 'Custom':
from custom import Custom
model = Custom(anchors=cfg['anchors'])
else:
model = models.__dict__[args.arch](anchors=cfg['anchors'])
if args.resume:
assert isfile(args.resume), '{} is not a valid file'.format(args.resume)
model = load_pretrain(model, args.resume)
model.eval()
model = model.to(device)
default_hp = cfg.get('hp', {})
p = dict()
p['network'] = model
p['network_name'] = args.arch+'_'+args.resume.split('/')[-1].split('.')[0]
p['dataset'] = args.dataset
global ims, gt, image_files
dataset_info = load_dataset(args.dataset)
videos = list(dataset_info.keys())
np.random.shuffle(videos)
for video in videos:
print(video)
if isfile('finish.flag'):
return
p['video'] = video
ims = None
image_files = dataset_info[video]['image_files']
gt = dataset_info[video]['gt']
np.random.shuffle(params['penalty_k'])
np.random.shuffle(params['window_influence'])
np.random.shuffle(params['lr'])
for penalty_k in params['penalty_k']:
for window_influence in params['window_influence']:
for lr in params['lr']:
for instance_size in params['instance_size']:
p['hp'] = default_hp.copy()
p['hp'].update({'penalty_k':penalty_k,
'window_influence':window_influence,
'lr':lr,
'instance_size': instance_size,
})
tune(p)
def main():
"""
基础网络的训练
:return:
"""
global args, best_acc, tb_writer, logger
args = parser.parse_args()
# 初始化日志信息
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
# 获取log信息
logger = logging.getLogger('global')
logger.info("\n" + collect_env_info())
logger.info(args)
# 获取配置信息
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# 构建数据集
train_loader, val_loader = build_data_loader(cfg)
# 加载训练网络
if args.arch == 'Custom':
from custom import Custom
model = Custom(pretrain=True, anchors=cfg['anchors'])
else:
exit()
logger.info(model)
# 加载预训练网络
if args.pretrained:
model = load_pretrain(model, args.pretrained)
# GPU版本
# model = model.cuda()
# dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()
# 网络模型
dist_model = torch.nn.DataParallel(model)
# 模型参数的更新比例
if args.resume and args.start_epoch != 0:
model.features.unfix((args.start_epoch - 1) / args.epochs)
# 获取优化器和学习率的更新策略
optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch)
# optionally resume from a checkpoint 加载模型
if args.resume:
assert os.path.isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model, optimizer, args.start_epoch, best_acc, arch = restore_from(
model, optimizer, args.resume)
# GPU
# dist_model = torch.nn.DataParallel(model, list(range(torch.cuda.device_count()))).cuda()
dist_model = torch.nn.DataParallel(model)
logger.info(lr_scheduler)
logger.info('model prepare done')
# 模型训练
train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch,
cfg)
def main():
global args, best_acc, tb_writer, logger
args = parser.parse_args()
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log, logging.INFO)
logger = logging.getLogger('global')
logger.info("\n" + collect_env_info())
logger.info(args)
cfg = load_config(args)
logger.info("config \n{}".format(json.dumps(cfg, indent=4)))
if args.log_dir:
tb_writer = SummaryWriter(args.log_dir)
else:
tb_writer = Dummy()
# build dataset
train_loader, val_loader = build_data_loader(cfg)
if args.arch == 'Custom':
model = Custom(anchors=cfg['anchors'])
elif args.arch == 'Custom_Sky':
model = Custom_Sky(anchors=cfg['anchors'])
else:
exit()
logger.info(model)
if args.pretrained:
model = load_pretrain(model, args.pretrained)
# print(summary(model=model, input_size=(3, 511, 511), batch_size=1))
model = model.cuda()
dist_model = torch.nn.DataParallel(model,
list(range(
torch.cuda.device_count()))).cuda()
if args.resume and args.start_epoch != 0:
model.features.unfix((args.start_epoch - 1) / args.epochs)
optimizer, lr_scheduler = build_opt_lr(model, cfg, args, args.start_epoch)
# optionally resume from a checkpoint
if args.resume:
print(args.resume)
assert os.path.isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model, optimizer, args.start_epoch, best_acc, arch = restore_from(
model, optimizer, args.resume)
dist_model = torch.nn.DataParallel(
model, list(range(torch.cuda.device_count()))).cuda()
logger.info(lr_scheduler)
logger.info('model prepare done')
train(train_loader, dist_model, optimizer, lr_scheduler, args.start_epoch,
cfg)
def main():
global args, logger, v_id #全局变量
args = parser.parse_args() #args是test.py文件运行时，接受的参数
cfg = load_config(args) #加载 JSON 配置文件并设置args.arch的值。
print(cfg)
init_log('global', logging.INFO)
if args.log != "":
add_file_handler('global', args.log,
logging.INFO) #add_file_handler 创建一个记录器并绑定文件句柄。
logger = logging.getLogger('global')
logger.info(args)
# setup model Custom 为论文实现的网络。如果不是“Custom”，加载 models 下指定的结构。
if args.arch == 'Custom': #args.arch参数，预训练模型的结构，命令行不给的话，默认为' ',
from custom import Custom
model = Custom(anchors=cfg['anchors']
) #cfg是从config_vot.json的到的数据，所以跟踪时用的model.anchors字典中的数据
else:
parser.error('invalid architecture: {}'.format(args.arch))
if args.resume: #给了args.resume,如果args.resume不是文件，报错，
assert isfile(args.resume), '{} is not a valid file'.format(
args.resume)
model = load_pretrain(
model, args.resume) #args.resume是文件load_pretrain ,能够处理网络之间的不一致
model.eval()
device = torch.device('cuda' if (
torch.cuda.is_available() and not args.cpu) else 'cpu')
model = model.to(device)
# setup dataset，字典
dataset = load_dataset(
args.dataset) #load_dataset 能够加载 VOT、DAVIS、ytb_vos 三种数据集。
#仅以上三种数据源支持掩膜输出。
# VOS or VOT?
if args.dataset in ['DAVIS2016', 'DAVIS2017', 'ytb_vos'] and args.mask:
vos_enable = True # enable Mask output ,使用掩膜输出
else:
vos_enable = False
total_lost = 0 # VOT 跟踪任务有损失函数
iou_lists = [] # VOS 分割任务
speed_list = []
#v_id视频索引从1起,video是视频名字
for v_id, video in enumerate(dataset.keys(), start=1):
if v_id == 2:
exit()
if args.video != '' and video != args.video: #不成立，args.video默认是' '
continue
if vos_enable: #分割任务,,,,分割任务和跟踪任务只能选一个
iou_list, speed = track_vos(
model,
dataset[video],
cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask,
args.refine,
args.dataset in ['DAVIS2017', 'ytb_vos'],
device=device)
iou_lists.append(iou_list) #iou_list是什么类型的数据？？？
else: #跟踪任务
lost, speed = track_vot(model,
dataset[video],
cfg['hp'] if 'hp' in cfg.keys() else None,
args.mask,
args.refine,
device=device)
total_lost += lost
speed_list.append(speed)
# report final result记录最终结果
if vos_enable: #如果进行的是分割任务
for thr, iou in zip(thrs, np.mean(np.concatenate(iou_lists), axis=0)):
logger.info('Segmentation Threshold {:.2f} mIoU: {:.3f}'.format(
thr, iou))
else:
logger.info('Total Lost: {:d}'.format(total_lost))
logger.info('Mean Speed: {:.2f} FPS'.format(np.mean(speed_list)))
