def write_predictions(self, output_dir=None):
if output_dir is None:
output_dir = self.output_dir
results = self.get_results()
if self.tracktor['interpolate']:
results = interpolate(results)
print(f"Writing predictions to: {output_dir}")
write_results(results, output_dir)
def get_res(img,
inference,
scale,
tracker,
depth_merger='mean',
given_K=False):
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_pil = Image.fromarray(img)
masks, masks_im, boxes = None, None, None
results = None
with torch.no_grad():
if inference['name'] == 'monodepth':
depth_map, depth_im = monodepth_infer(inference['encoder'],
inference['depth_decoder'],
inference['input_size'],
img_pil)
depth_map = depth_map[0, 0] # * 5.4
else: # mannequin
depth_map, depth_im = mannequin_infer(img_pil)
depth_map = (255 - depth_map) / 7
if tracker == None:
masks, masks_im, boxes = infer_segmentation(
"yolact_plus_resnet50_54_800000.pth", img)
else:
transforms = Compose([
Resize((749, 1333)),
ToTensor(),
])
frame_batch = {
'img': transforms(img_pil).unsqueeze(0) #.cuda()
}
tracker.step(frame_batch)
results = tracker.get_results()
results = interpolate(results)
res_img = img.copy()
if tracker == None:
return img, merge_masks(res_img, masks, masks_im, boxes, depth_merger,
depth_map, inference['name'], scale, given_K)
return img, merge_boxes(res_img, results, depth_merger, depth_map,
inference['name'])
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# neural motion model
vis_model = VisSimpleReID()
motion_model = MotionModelV3(vis_model)
motion_model.load_state_dict(torch.load('output/motion/finetune_motion_model_v3.pth'))
motion_model.eval()
motion_model.cuda()
save_vis_results = False
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle'])
else:
# tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
tracker = TrackerNeuralMM(obj_detect, reid_network, motion_model, tracktor['tracker'], save_vis_results=save_vis_results, vis_model=None)
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'], {'use_val_split':True})
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(seq) * tracktor['frame_split'][1]:
with torch.no_grad():
tracker.step(frame)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if save_vis_results:
vis_results = tracker.get_vis_results()
seq.write_vis_results(vis_results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq, osp.join(output_dir, tracktor['dataset'], str(seq)))
_log.info(f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'], tracktor['output_subdir'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2).to(device)
obj_detect.load_state_dict(
torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn']).to(device)
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
tracker.step(frame, i)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
summary = evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
summary.to_pickle(
"output/finetuning_results/results_{}_{}_{}_{}_{}.pkl".format(
tracktor['output_subdir'],
tracktor['tracker']['finetuning']['max_displacement'],
tracktor['tracker']['finetuning']['batch_size'],
tracktor['tracker']['finetuning']['learning_rate'],
tracktor['tracker']['finetuning']['iterations']))
def my_main(tracktor, siamese, _config):
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
print("[*] Building object detector")
if tracktor['network'].startswith('frcnn'):
# FRCNN
from tracktor.frcnn import FRCNN
from frcnn.model import config
if _config['frcnn']['cfg_file']:
config.cfg_from_file(_config['frcnn']['cfg_file'])
if _config['frcnn']['set_cfgs']:
config.cfg_from_list(_config['frcnn']['set_cfgs'])
obj_detect = FRCNN(num_layers=101)
obj_detect.create_architecture(2,
tag='default',
anchor_scales=config.cfg.ANCHOR_SCALES,
anchor_ratios=config.cfg.ANCHOR_RATIOS)
obj_detect.load_state_dict(torch.load(tracktor['obj_detect_weights']))
elif tracktor['network'].startswith('fpn'):
# FPN
from tracktor.fpn import FPN
from fpn.model.utils import config
config.cfg.TRAIN.USE_FLIPPED = False
config.cfg.CUDA = True
config.cfg.TRAIN.USE_FLIPPED = False
checkpoint = torch.load(tracktor['obj_detect_weights'])
if 'pooling_mode' in checkpoint.keys():
config.cfg.POOLING_MODE = checkpoint['pooling_mode']
set_cfgs = [
'ANCHOR_SCALES', '[4, 8, 16, 32]', 'ANCHOR_RATIOS', '[0.5,1,2]'
]
config.cfg_from_file(_config['tracktor']['obj_detect_config'])
config.cfg_from_list(set_cfgs)
obj_detect = FPN(('__background__', 'pedestrian'),
101,
pretrained=False)
obj_detect.create_architecture()
obj_detect.load_state_dict(checkpoint['model'])
else:
raise NotImplementedError(
f"Object detector type not known: {tracktor['network']}")
pprint.pprint(config.cfg)
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **siamese['cnn'])
reid_network.load_state_dict(torch.load(tracktor['reid_network_weights']))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
print("[*] Beginning evaluation...")
time_total = 0
for sequence in Datasets(tracktor['dataset']):
tracker.reset()
now = time.time()
print("[*] Evaluating: {}".format(sequence))
data_loader = DataLoader(sequence, batch_size=1, shuffle=False)
for i, frame in enumerate(data_loader):
# frame_split = [0.0, 1.0]
if i >= len(sequence) * tracktor['frame_split'][0] and i <= len(
sequence) * tracktor['frame_split'][1]:
tracker.step(frame)
results = tracker.get_results()
time_total += time.time() - now
print("[*] Tracks found: {}".format(len(results)))
print("[*] Time needed for {} evaluation: {:.3f} s".format(
sequence,
time.time() - now))
if tracktor['interpolate']:
results = interpolate(results)
sequence.write_results(results, osp.join(output_dir))
if tracktor['write_images']:
plot_sequence(
results, sequence,
osp.join(output_dir, tracktor['dataset'], str(sequence)))
print("[*] Evaluation for all sets (without image generation): {:.3f} s".
format(time_total))
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# motion network
motion_network = None
if tracktor['tracker']['motion_model_enabled'] and not tracktor['motion'][
'use_cva_model']:
motion_network = eval(
tracktor['motion']['model'])(**tracktor['motion']['model_args'])
motion_network.load_state_dict(
torch.load(tracktor['motion']['network_weights'])['model'])
motion_network.eval().cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, motion_network,
tracktor['tracker'], tracktor['motion'], 2)
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
start = time.time()
all_mm_times = []
all_warp_times = []
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
with torch.no_grad():
mm_time, warp_time = tracker.step(frame)
if mm_time is not None:
all_mm_times.append(mm_time)
if warp_time is not None:
all_warp_times.append(warp_time)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
_log.info(
f"Average FPS for {seq}: {len(data_loader) / (time.time() - start) :.3f}"
)
_log.info(
f"Average MM time for {seq}: {float(np.array(all_mm_times).mean()) :.3f} s"
)
if all_warp_times:
_log.info(
f"Average warp time for {seq}: {float(np.array(all_warp_times).mean()) :.3f} s"
)
if tracktor['interpolate']:
results = interpolate(results)
if 'semi_online' in tracktor and tracktor['semi_online']:
for i, track in results.items():
for frame in sorted(track, reverse=True):
if track[frame][5] == 0:
break
del track[frame]
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)),
tracktor['tracker']['plot_mm'])
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)"
)
if mot_accums:
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
def main(tracktor, reid, _config, _log, _run):
target = Target()
targetpath = target.Folder()
targetname = target.TargetName()
vottpath = target.GetVottPath()
vottfile = target.GetVottContent()
dictid, timelist = target.GetTagTime(vottfile)
print(f"{len(timelist)} frames were tagged")
timedict = target.ExtractByTimeList(timelist)
bbdict = target.GetbbWithTime(vottfile)
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
print("Tracktor初始化完成")
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
print(f"{seq}加載完成, tracking開始")
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
id = tracker.step(frame, bbdict[timedict["%06d" % num_frames]])
target.WriteID2asset(id, dictid[timedict["%06d" % num_frames]])
num_frames += 1
results = tracker.get_results()
ids = list(results.keys())
target.WriteID2vott(ids, vottfile=vottfile)
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
target.CleanImg()
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
if tracktor['write_videos']:
plot_sequence_video(
results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
for i, frame in enumerate(tqdm(dataset)):
if len(dataset) * tracktor['frame_split'][0] <= i <= len(
dataset) * tracktor['frame_split'][1]:
with torch.no_grad():
tracker.step(frame)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
print(f"Tracks found: {len(results)}")
print(f"Runtime for {dataset}: {time.time() - start :.2f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if dataset.no_gt:
print(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, dataset))
print(f"Writing predictions to: {output_dir}")
dataset.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, dataset, osp.join(output_dir))
print(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)"
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
# Data transform
normalize_mean=[0.485, 0.456, 0.406]
normalize_std=[0.229, 0.224, 0.225]
# dataset = Datasets(tracktor['dataset'])
transforms = ToTensor()
# transforms = Compose([ToTensor(), Normalize(normalize_mean,
# normalize_std)])
tracker.reset()
# tracker.public_detections=False
start = time.time()
_log.info(f"Tracking: video")
# Load video and annotations
cap = cv2.VideoCapture("/home/yc3390/camera_detection_demo/data/prid2011_videos/test_b_1min_1min.mp4")
with open("/home/yc3390/camera_detection_demo/data/prid2011_videos/anno_b.pkl", 'rb') as f:
gts = pk.load(f)
det_file = "/data/yc3390/tracktor_output/output/tracktor/MOT17/Tracktor++/Video-result_ReID.txt"
# with open("/data/yc3390/tracktor_output/output/tracktor/MOT17/Tracktor++/Video-result_ReID.pkl", 'rb') as f:
# dts = pk.load(f)
# for dt in dts:
# if len(dt['boxes'][0]):
# for i in range(len(dt['boxes'])):
# dt['boxes'][i][-1] = -1
offset = 25 * 60
dets = {}
for i in range(1, offset+1):
dets[i] = []
assert osp.exists(det_file)
with open(det_file, "r") as inf:
reader = csv.reader(inf, delimiter=',')
for row in reader:
x1 = float(row[2]) - 1
y1 = float(row[3]) - 1
# This -1 accounts for the width (width of 1 x1=x2)
x2 = x1 + float(row[4]) - 1
y2 = y1 + float(row[5]) - 1
score = float(row[6])
bb = np.array([x1,y1,x2,y2], dtype=np.float32)
dets[int(float(row[0]))].append(bb)
frame_count = offset
while True:
ret, image = cap.read()
if not ret:
break
# BGR to RGB
image = Image.fromarray(image[..., ::-1])
image = transforms(image)[None, ...]
# Detection
# if frame_count in gts.keys():
# frames =
blob = {"dets" : torch.Tensor([dets[i]]), "img" : image}
tracker.step(blob)
frame_count += 1
print("Finished ", frame_count, output_dir, image.shape)
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for video: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if True:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
write_results(results, output_dir)
def main():
args = parse_args()
logging.basicConfig(level=logging.INFO,format='%(asctime)s %(name)-12s %(levelname)-8s %(message)s')
logging.root.setLevel(logging.INFO)
is_cuda = torch.cuda.is_available()
if is_cuda:
LOG.info('-' * 50)
LOG.info('Enabling CUDA')
LOG.info('-' * 50)
device = torch.device('cuda' if is_cuda else 'cpu')
# sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
if is_cuda:
torch.cuda.manual_seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name'])
if not osp.exists(output_dir):
os.makedirs(output_dir)
##########################
# Initialize the modules #
##########################
# object detection
LOG.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect_state_dict = torch.load(args.detection_path, map_location=device)
obj_detect.load_state_dict(obj_detect_state_dict)
obj_detect.eval()
if is_cuda:
obj_detect.cuda()
# LOG.info('Load detection model...')
# obj_detect = load_object_detection_driver(args.detection_path)
# LOG.info('Done.')
# reid
LOG.info("Initializing reidentification network.")
reid_network = resnet50(pretrained=False, output_dim=128)
reid_network.load_state_dict(torch.load(args.reid_path, map_location=device))
reid_network.eval()
if is_cuda:
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'], tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
tracker.reset()
start = time.time()
vc = cv2.VideoCapture(args.source)
frame_count = vc.get(cv2.CAP_PROP_FRAME_COUNT)
video_output = args.output
fourcc = cv2.VideoWriter_fourcc(*'avc1')
width = int(vc.get(cv2.CAP_PROP_FRAME_WIDTH))
height = int(vc.get(cv2.CAP_PROP_FRAME_HEIGHT))
fps = vc.get(cv2.CAP_PROP_FPS)
each_frame = args.every_frame
writer = cv2.VideoWriter(video_output, fourcc, fps / each_frame, frameSize=(width, height))
LOG.info(f"Tracking: {args.source}")
frame_id = 0
frame_num = 0
results = {}
try:
while True:
frame_num += 1
if frame_num % each_frame == 0:
ret, frame = vc.read()
if not ret:
break
else:
vc.grab()
continue
frame_id += 1
if frame_id % 50 == 0:
LOG.info(f'Processing frame {frame_id}')
if frame_count * tracktor['frame_split'][0] <= frame_id <= frame_count * tracktor['frame_split'][1]:
rgb_frame = frame[:, :, ::-1]
torch_frame = F.to_tensor(rgb_frame.copy())
torch_frame = torch_frame.expand([1, *torch_frame.shape])
if is_cuda:
torch_frame = torch_frame.cuda()
torch_blob = {
'img': torch_frame
}
tracker.step(torch_blob, frame)
# __import__('ipdb').set_trace()
results = tracker.results
output = draw_boxes(frame, frame_id - 1, results=results)
writer.write(output)
except KeyboardInterrupt:
LOG.info('Stopping.')
writer.release()
time_total += time.time() - start
LOG.info(f"Tracks found: {len(results)}")
LOG.info(f"Runtime for {args.source}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = utils.interpolate(results)
# if tracktor['write_images']:
# utils.plot_sequence(results, seq, osp.join(output_dir, args.source))
LOG.info(f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({frame_id / time_total:.1f} Hz)")
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
_log.info("Initializing object detector.")
# object detection
obj_detect = FRCNN_FPN(num_classes=2, correlation_head=CorrelationHead())
obj_detect_model = torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage)
correlation_weights = torch.load(
_config['tracktor']['correlation_weights'],
map_location=lambda storage, loc: storage)
for k in correlation_weights:
obj_detect_model.update(
{"correlation_head." + k: correlation_weights[k]})
obj_detect.load_state_dict(obj_detect_model)
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
with torch.no_grad():
tracker.step(frame)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.2f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
score_killed_tracks = tracker.get_score_killed_tracks()
_log.info(f"Score Killed Tracks: ({len(score_killed_tracks)})")
for kill in score_killed_tracks:
_log.info(
f"Track [ {kill['id']:3d} ] killed in frame [ {kill['frame']:3d} ]"
)
nms_killed_tracks = tracker.get_nms_killed_tracks()
_log.info(f"NMS Killed Tracks ({len(nms_killed_tracks)}):")
for kill in nms_killed_tracks:
_log.info(
f"Track [ {kill['id']:3d} ] killed in frame [ {kill['frame']:3d} ]"
)
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.2f} s for {num_frames} frames ({num_frames / time_total:.2f} Hz)"
)
if mot_accums:
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']),
tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
obj_detect = FRCNN_FPN(num_classes=2)
obj_detect.load_state_dict(
torch.load(_config['tracktor']['obj_detect_model'],
map_location=lambda storage, loc: storage))
obj_detect.eval()
obj_detect.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(
torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(obj_detect, reid_network, tracktor['tracker'],
tracktor['oracle'])
else:
tracker = Tracker(obj_detect, reid_network, tracktor['tracker'])
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(
seq) * tracktor['frame_split'][1]:
tracker.step(frame)
num_frames += 1
results = tracker.get_results()
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
if tracktor['write_images']:
plot_sequence(results, seq,
osp.join(output_dir, tracktor['dataset'], str(seq)))
img_array = []
dir = osp.join(output_dir, tracktor['dataset'], str(seq), "*.jpg")
files = glob.glob(dir)
sorted_files = natsorted(files)
for filename in sorted_files:
img = cv2.imread(filename)
height, width, layers = img.shape
size = (width, height)
img_array.append(img)
out = cv2.VideoWriter(
osp.join(output_dir, tracktor['dataset'],
str(seq), "result_video.avi"),
cv2.VideoWriter_fourcc(*'DIVX'), 10, size)
for i in range(len(img_array)):
out.write(img_array[i])
out.release()
_log.info(
f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
evaluate_mot_accums(mot_accums,
[str(s) for s in dataset if not s.no_gt],
generate_overall=True)
def main(tracktor, reid, _config, _log, _run):
sacred.commands.print_config(_run)
# set all seeds
torch.manual_seed(tracktor['seed'])
torch.cuda.manual_seed(tracktor['seed'])
np.random.seed(tracktor['seed'])
torch.backends.cudnn.deterministic = True
output_dir = osp.join(get_output_dir(tracktor['module_name']), tracktor['name'])
sacred_config = osp.join(output_dir, 'sacred_config.yaml')
if not osp.exists(output_dir):
os.makedirs(output_dir)
with open(sacred_config, 'w') as outfile:
yaml.dump(_config, outfile, default_flow_style=False)
##########################
# Initialize the modules #
##########################
# object detection
_log.info("Initializing object detector.")
use_masks = _config['tracktor']['tracker']['use_masks']
mask_model = Mask_RCNN(num_classes=2)
fast_model = FRCNN_FPN(num_classes=2)
fast_model.load_state_dict(torch.load(_config['tracktor']['fast_rcnn_model'],
map_location=lambda storage, loc: storage))
if(use_masks):
mask_model.load_state_dict(torch.load(_config['tracktor']['mask_rcnn_model'],
map_location=lambda storage, loc: storage)['model_state_dict'])
mask_model.eval()
mask_model.cuda()
fast_model.eval()
fast_model.cuda()
# reid
reid_network = resnet50(pretrained=False, **reid['cnn'])
reid_network.load_state_dict(torch.load(tracktor['reid_weights'],
map_location=lambda storage, loc: storage))
reid_network.eval()
reid_network.cuda()
# tracktor
if 'oracle' in tracktor:
tracker = OracleTracker(fast_model, reid_network, tracktor['tracker'], tracktor['oracle'])
else:
tracker = Tracker(fast_model, reid_network, tracktor['tracker'], mask_model)
time_total = 0
num_frames = 0
mot_accums = []
dataset = Datasets(tracktor['dataset'])
for seq in dataset:
num_frames = 0
tracker.reset()
start = time.time()
_log.info(f"Tracking: {seq}")
data_loader = DataLoader(seq, batch_size=1, shuffle=False)
if tracktor['write_images'] and use_masks:
print("[*] Plotting image to {}".format(osp.join(output_dir, tracktor['dataset'])))
for i, frame in enumerate(tqdm(data_loader)):
if len(seq) * tracktor['frame_split'][0] <= i <= len(seq) * tracktor['frame_split'][1]:
tracker.step(frame)
if tracktor['write_images'] and use_masks:
result = tracker.get_results()
masks = tracker.get_masks()
plot_sequence(result, masks, seq, num_frames, osp.join(output_dir, tracktor['dataset'], str(seq)), plot_masks = True)
num_frames += 1
results = tracker.get_results()
import matplotlib.pyplot as plt
time_total += time.time() - start
_log.info(f"Tracks found: {len(results)}")
_log.info(f"Runtime for {seq}: {time.time() - start :.1f} s.")
if tracktor['interpolate']:
results = interpolate(results)
if seq.no_gt:
_log.info(f"No GT data for evaluation available.")
else:
mot_accums.append(get_mot_accum(results, seq))
_log.info(f"Writing predictions to: {output_dir}")
seq.write_results(results, output_dir)
_log.info(f"Tracking runtime for all sequences (without evaluation or image writing): "
f"{time_total:.1f} s ({num_frames / time_total:.1f} Hz)")
if mot_accums:
evaluate_mot_accums(mot_accums, [str(s) for s in dataset if not s.no_gt], generate_overall=True)
