def __init__(self,
update_facebank=False,
show_results=True,
dataset_path=None,
save=False):
"""
:param update_facebank: Calculate new facebank or use old one
:param show_results: Draw bboxes and ident info at frame
:param videoPath: Path for processing the video
:param save: Save processed video
"""
self.dataset_path = dataset_path
self.video_path = f'{self.dataset_path}/{self.find_video()}'
self.face_identificator = FaceIdentificator(
update_facebank, f'{os.getcwd()}/{self.dataset_path}')
self.tracker = Tracker()
self.unique_facebank_names = list(
self.face_identificator.saved_embeddings.keys())
# TODO Create smth more smarter for scores initialization
self.saved_scores = dict(
zip(self.unique_facebank_names, [
self.face_identificator.threshold
for name in self.unique_facebank_names
]))
self.mapped_tracks = {}
self.show_results = show_results
self.treshold = 0.4
self.save = save
def main(args):
utils.init_distributed_mode(args)
print("git:\n {}\n".format(utils.get_sha()))
if args.frozen_weights is not None:
assert args.masks, "Frozen training is meant for segmentation only"
print(args)
device = torch.device(args.device)
# fix the seed for reproducibility
seed = args.seed + utils.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
scaler = torch.cuda.amp.GradScaler(enabled=args.fp16)
if args.det_val:
assert args.eval, 'only support eval mode of detector for track'
model, criterion, postprocessors = build_model(args)
elif args.eval:
model, criterion, postprocessors = build_tracktest_model(args)
else:
model, criterion, postprocessors = build_tracktrain_model(args)
model.to(device)
model_without_ddp = model
n_parameters = sum(p.numel() for p in model.parameters()
if p.requires_grad)
print('number of params:', n_parameters)
dataset_train = build_dataset(image_set=args.track_train_split, args=args)
dataset_val = build_dataset(image_set=args.track_eval_split, args=args)
if args.distributed:
if args.cache_mode:
sampler_train = samplers.NodeDistributedSampler(dataset_train)
sampler_val = samplers.NodeDistributedSampler(dataset_val,
shuffle=False)
else:
sampler_train = samplers.DistributedSampler(dataset_train)
sampler_val = samplers.DistributedSampler(dataset_val,
shuffle=False)
else:
sampler_train = torch.utils.data.RandomSampler(dataset_train)
sampler_val = torch.utils.data.SequentialSampler(dataset_val)
batch_sampler_train = torch.utils.data.BatchSampler(sampler_train,
args.batch_size,
drop_last=True)
data_loader_train = DataLoader(dataset_train,
batch_sampler=batch_sampler_train,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
data_loader_val = DataLoader(dataset_val,
args.batch_size,
sampler=sampler_val,
drop_last=False,
collate_fn=utils.collate_fn,
num_workers=args.num_workers,
pin_memory=True)
# lr_backbone_names = ["backbone.0", "backbone.neck", "input_proj", "transformer.encoder"]
def match_name_keywords(n, name_keywords):
out = False
for b in name_keywords:
if b in n:
out = True
break
return out
for n, p in model_without_ddp.named_parameters():
print(n)
param_dicts = [{
"params": [
p for n, p in model_without_ddp.named_parameters()
if not match_name_keywords(n, args.lr_backbone_names)
and not match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters() if
match_name_keywords(n, args.lr_backbone_names) and p.requires_grad
],
"lr":
args.lr_backbone,
}, {
"params": [
p for n, p in model_without_ddp.named_parameters()
if match_name_keywords(n, args.lr_linear_proj_names)
and p.requires_grad
],
"lr":
args.lr * args.lr_linear_proj_mult,
}]
if args.sgd:
optimizer = torch.optim.SGD(param_dicts,
lr=args.lr,
momentum=0.9,
weight_decay=args.weight_decay)
else:
optimizer = torch.optim.AdamW(param_dicts,
lr=args.lr,
weight_decay=args.weight_decay)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, args.lr_drop)
if args.distributed:
model = torch.nn.parallel.DistributedDataParallel(
model, device_ids=[args.gpu], find_unused_parameters=True)
model_without_ddp = model.module
if args.dataset_file == "coco_panoptic":
# We also evaluate AP during panoptic training, on original coco DS
coco_val = datasets.coco.build("val", args)
base_ds = get_coco_api_from_dataset(coco_val)
else:
base_ds = get_coco_api_from_dataset(dataset_val)
if args.frozen_weights is not None:
checkpoint = torch.load(args.frozen_weights, map_location='cpu')
model_without_ddp.detr.load_state_dict(checkpoint['model'])
output_dir = Path(args.output_dir)
if args.resume:
if args.resume.startswith('https'):
checkpoint = torch.hub.load_state_dict_from_url(args.resume,
map_location='cpu',
check_hash=True)
else:
checkpoint = torch.load(args.resume, map_location='cpu')
missing_keys, unexpected_keys = model_without_ddp.load_state_dict(
checkpoint['model'], strict=False)
unexpected_keys = [
k for k in unexpected_keys
if not (k.endswith('total_params') or k.endswith('total_ops'))
]
if len(missing_keys) > 0:
print('Missing Keys: {}'.format(missing_keys))
if len(unexpected_keys) > 0:
print('Unexpected Keys: {}'.format(unexpected_keys))
if not args.eval and 'optimizer' in checkpoint and 'lr_scheduler' in checkpoint and 'epoch' in checkpoint:
import copy
p_groups = copy.deepcopy(optimizer.param_groups)
optimizer.load_state_dict(checkpoint['optimizer'])
for pg, pg_old in zip(optimizer.param_groups, p_groups):
pg['lr'] = pg_old['lr']
pg['initial_lr'] = pg_old['initial_lr']
print(optimizer.param_groups)
lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
# todo: this is a hack for doing experiment that resume from checkpoint and also modify lr scheduler (e.g., decrease lr in advance).
args.override_resumed_lr_drop = True
if args.override_resumed_lr_drop:
print(
'Warning: (hack) args.override_resumed_lr_drop is set to True, so args.lr_drop would override lr_drop in resumed lr_scheduler.'
)
lr_scheduler.step_size = args.lr_drop
lr_scheduler.base_lrs = list(
map(lambda group: group['initial_lr'],
optimizer.param_groups))
lr_scheduler.step(lr_scheduler.last_epoch)
args.start_epoch = checkpoint['epoch'] + 1
# check the resumed model
# if not args.eval:
# test_stats, coco_evaluator, _ = evaluate(
# model, criterion, postprocessors, data_loader_val, base_ds, device, args.output_dir
# )
if args.eval:
assert args.batch_size == 1, print("Now only support 1.")
tracker = Tracker(score_thresh=args.track_thresh)
test_stats, coco_evaluator, res_tracks = evaluate(model,
criterion,
postprocessors,
data_loader_val,
base_ds,
device,
args.output_dir,
tracker=tracker,
phase='eval',
det_val=args.det_val,
fp16=args.fp16)
if args.output_dir:
# utils.save_on_master(coco_evaluator.coco_eval["bbox"].eval, output_dir / "eval.pth")
if res_tracks is not None:
print("Creating video index for {}.".format(args.dataset_file))
video_to_images = defaultdict(list)
video_names = defaultdict()
for _, info in dataset_val.coco.imgs.items():
video_to_images[info["video_id"]].append({
"image_id":
info["id"],
"frame_id":
info["frame_id"]
})
video_name = info["file_name"].split("/")[0]
if video_name not in video_names:
video_names[info["video_id"]] = video_name
assert len(video_to_images) == len(video_names)
# save mot results.
save_track(res_tracks, args.output_dir, video_to_images,
video_names, args.track_eval_split)
return
print("Start training")
start_time = time.time()
for epoch in range(args.start_epoch, args.epochs):
if args.distributed:
sampler_train.set_epoch(epoch)
train_stats = train_one_epoch(model,
criterion,
data_loader_train,
optimizer,
device,
scaler,
epoch,
args.clip_max_norm,
fp16=args.fp16)
lr_scheduler.step()
if args.output_dir:
checkpoint_paths = [output_dir / 'checkpoint.pth']
# extra checkpoint before LR drop and every 5 epochs
if (epoch + 1) % args.lr_drop == 0 or (epoch + 1) % 5 == 0:
checkpoint_paths.append(output_dir /
f'checkpoint{epoch:04}.pth')
for checkpoint_path in checkpoint_paths:
utils.save_on_master(
{
'model': model_without_ddp.state_dict(),
'optimizer': optimizer.state_dict(),
'lr_scheduler': lr_scheduler.state_dict(),
'epoch': epoch,
'args': args,
}, checkpoint_path)
log_stats = {
**{f'train_{k}': v
for k, v in train_stats.items()}, 'epoch': epoch,
'n_parameters': n_parameters
}
if epoch % 10 == 0 or epoch > args.epochs - 5:
test_stats, coco_evaluator, _ = evaluate(model,
criterion,
postprocessors,
data_loader_val,
base_ds,
device,
args.output_dir,
fp16=args.fp16)
log_test_stats = {
**{f'test_{k}': v
for k, v in test_stats.items()}
}
log_stats.update(log_test_stats)
if args.output_dir and utils.is_main_process():
with (output_dir / "log.txt").open("a") as f:
f.write(json.dumps(log_stats) + "\n")
# for evaluation logs
# if coco_evaluator is not None:
# (output_dir / 'eval').mkdir(exist_ok=True)
# if "bbox" in coco_evaluator.coco_eval:
# filenames = ['latest.pth']
# if epoch % 50 == 0:
# filenames.append(f'{epoch:03}.pth')
# for name in filenames:
# torch.save(coco_evaluator.coco_eval["bbox"].eval,
# output_dir / "eval" / name)
total_time = time.time() - start_time
total_time_str = str(datetime.timedelta(seconds=int(total_time)))
print('Training time {}'.format(total_time_str))
def main(args):
cap = cv2.VideoCapture(args.video_input)
frame_rate = int(round(cap.get(cv2.CAP_PROP_FPS)))
video_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
video_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
demo_images_path = os.path.join(args.demo_output, 'demo_images')
if not os.path.exists(demo_images_path):
os.makedirs(demo_images_path)
device = torch.device(args.device)
model, _, postprocessors = build_tracktest_model(args)
model.to(device)
model.eval()
tracker = Tracker(score_thresh=args.track_thresh)
checkpoint = torch.load(args.resume, map_location='cpu')
_, _ = model.load_state_dict(checkpoint['model'], strict=False)
print("Model is loaded")
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
color_list = colormap()
print("Starting inference")
count = 0
tracker.reset_all()
pre_embed = None
res, img = cap.read()
while res:
count += 1
resized_img, nh, nw = resize(img)
rgb_img = cv2.cvtColor(resized_img, cv2.COLOR_BGR2RGB)
tensor_img = F.normalize(F.to_tensor(rgb_img), mean, std)
samples = nested_tensor_from_tensor_list([tensor_img]).to(device)
outputs, pre_embed = model(samples, pre_embed)
orig_sizes = torch.stack(
[torch.as_tensor([video_height, video_width])], dim=0).to(device)
results = postprocessors['bbox'](outputs, orig_sizes)
if count == 1:
res_track = tracker.init_track(results[0])
else:
res_track = tracker.step(results[0])
for ret in res_track:
if ret['active'] == 0:
continue
bbox = ret['bbox']
tracking_id = ret['tracking_id']
cv2.rectangle(img, (int(bbox[0]), int(bbox[1])),
(int(bbox[2]), int(bbox[3])),
color_list[tracking_id % 79].tolist(),
thickness=2)
cv2.putText(img, "{}".format(tracking_id),
(int(bbox[0]), int(bbox[1])), cv2.FONT_HERSHEY_SIMPLEX,
0.8, color_list[tracking_id % 79].tolist(), 2)
cv2.imwrite(
os.path.join(demo_images_path, "demo{:0>6d}.png".format(count)),
img)
print('Frame{:d} of the video is done'.format(count))
res, img = cap.read()
print('Lenth of the video: {:d} frames'.format(count))
print("Starting img2video")
img_paths = gb.glob(os.path.join(demo_images_path, "*.png"))
size = (video_width, video_height)
videowriter = cv2.VideoWriter(
os.path.join(args.demo_output, "demo_video.avi"),
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), frame_rate, size)
for img_path in sorted(img_paths):
img = cv2.imread(img_path)
img = cv2.resize(img, size)
videowriter.write(img)
videowriter.release()
print("img2video is done")
'init_beta':
1,
'init_v':
np.diag([4, 2]),
'init_nu':
5,
}
from models import GgiwTracker as Tracker
#cell 8
### Run Tracker
#cell 9
# tracker definition
tracker = Tracker(dt=dt, **config)
pr = cProfile.Profile()
for i in range(steps):
scan = measurements[measurements['ts'] == i]
pr.enable()
tracker.step(scan)
pr.disable()
estimates, log_lik = tracker.extract()
bboxes = tracker.extrackt_bbox()
s = io.StringIO()
ps = pstats.Stats(pr, stream=s).sort_stats('cumulative')
ps.print_stats(10)