def initialize(self, seq_info):
super(DAN, self).initialize(seq_info)
dataset_name = seq_info["dataset_name"]
seq_name = seq_info["seq_name"]
if dataset_name == "MOT15" and seq_name == "AVG-TownCentre":
self.choice = (4, 0, 4, 4, 5, 4)
init_test_mot15()
elif dataset_name == "MOT15":
self.choice = (0, 0, 4, 4, 5, 4)
init_test_mot15()
elif dataset_name == "DETRAC":
self.choice = TrackerConfig.get_ua_choice()
init_test_ua()
else:
self.choice = (0, 0, 4, 0, 3, 3)
init_test_mot17()
TrackerConfig.set_configure(self.choice)
if dataset_name == "MOT17":
self.min_confidence = 0.0
else:
self.min_confidence = None
config["resume"] = self.model_path
self.tracker = SSTTracker()
def create(c):
if not os.path.exists(image_folder) or not os.path.exists(
detection_file_name):
raise FileNotFoundError(
'cannot find the image folder and the detection file')
if not os.path.exists(args.log_folder):
os.mkdir(args.log_folder)
tracker = SSTTracker()
reader = MOTDataReader(image_folder=image_folder,
detection_file_name=detection_file_name,
min_confidence=0.0)
select_squences = [402, 404, 410, 422]
frame_index = 0
for i, item in enumerate(reader):
if i not in select_squences:
continue
if i > len(reader):
break
if item is None:
continue
img = item[0]
det = item[1]
if img is None or det is None or len(det) == 0:
continue
if len(det) > config['max_object']:
det = det[:config['max_object'], :]
h, w, _ = img.shape
det[:, [2, 4]] /= float(w)
det[:, [3, 5]] /= float(h)
image_org = tracker.update(img, det[:, 2:6], args.show_image,
frame_index)
frame_index += 1
if args.show_image and not image_org is None:
# det[:, [2, 4]] *= float(w)
# det[:, [3, 5]] *= float(h)
# boxes = det[:, 2:6].astype(int)
# for bid, b in enumerate(boxes):
# image_org = cv2.putText(image_org, '{}'.format(bid), tuple(b[:2]), cv2.FONT_HERSHEY_SIMPLEX, 1,
# (0, 0, 0), 2)
cv2.imshow('res', image_org)
cv2.imwrite(os.path.join(args.log_folder, '{0:06}.jpg'.format(i)),
image_org)
# cv2.waitKey(0)
print('frame: {}'.format(i))
def test(choice=None):
if args.type == 'train':
dataset_index = [2, 4, 5, 9, 10, 11, 13]
dataset_detection_type = {'-DPM', '-FRCNN', '-SDP'}
if args.type == 'test':
dataset_index = [1, 3, 6, 7, 8, 12, 14]
dataset_detection_type = {'-FRCNN', '-SDP', '-DPM'}
dataset_image_folder_format = os.path.join(
args.mot_root,
args.type + '/MOT' + str(args.mot_version) + '-{:02}{}/img1')
detection_file_name_format = os.path.join(
args.mot_root,
args.type + '/MOT' + str(args.mot_version) + '-{:02}{}/det/det.txt')
if not os.path.exists(args.log_folder):
os.mkdir(args.log_folder)
save_folder = ''
choice_str = ''
if not choice is None:
choice_str = TrackerConfig.get_configure_str(choice)
save_folder = os.path.join(args.log_folder, choice_str)
if not os.path.exists(save_folder):
os.mkdir(save_folder)
# else:
# return
saved_file_name_format = os.path.join(
save_folder, 'MOT' + str(args.mot_version) + '-{:02}{}.txt')
save_video_name_format = os.path.join(
save_folder, 'MOT' + str(args.mot_version) + '-{:02}{}.avi')
f = lambda format_str: [
format_str.format(index, type) for type in dataset_detection_type
for index in dataset_index
]
timer = Timer()
for image_folder, detection_file_name, saved_file_name, save_video_name in zip(
f(dataset_image_folder_format), f(detection_file_name_format),
f(saved_file_name_format), f(save_video_name_format)):
print('start processing ' + saved_file_name)
tracker = SSTTracker()
reader = MOTDataReader(image_folder=image_folder,
detection_file_name=detection_file_name,
min_confidence=0.0)
result = list()
result_str = saved_file_name
first_run = True
for i, item in enumerate(reader):
if i > len(reader):
break
if item is None:
continue
img = item[0]
det = item[1]
if img is None or det is None or len(det) == 0:
continue
if len(det) > config['max_object']:
det = det[:config['max_object'], :]
h, w, _ = img.shape
if first_run and args.save_video:
vw = cv2.VideoWriter(
save_video_name,
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 10, (w, h))
first_run = False
det[:, [2, 4]] /= float(w)
det[:, [3, 5]] /= float(h)
timer.tic()
image_org = tracker.update(img, det[:, 2:6], args.show_image, i)
timer.toc()
print('{}:{}, {}, {}\r'.format(os.path.basename(saved_file_name),
i, int(i * 100 / len(reader)),
choice_str))
#if args.show_image and not image_org is None:
#plt.imshow('res', image_org)
#cv2.waitKey(1)
if args.save_video and not image_org is None:
vw.write(image_org)
# save result
for t in tracker.tracks:
n = t.nodes[-1]
if t.age == 1:
b = n.get_box(tracker.frame_index - 1, tracker.recorder)
result.append([i] + [t.id] +
[b[0] * w, b[1] * h, b[2] * w, b[3] * h] +
[-1, -1, -1, -1])
# save data
np.savetxt(saved_file_name, np.array(result).astype(int), fmt='%i')
print(result_str)
print(timer.total_time)
print(timer.average_time)
def main(args):
base_net = SSTTracker()
def main(args):
np.random.seed(args.seed)
random.seed(args.seed)
torch.manual_seed(args.seed)
torch.cuda.manual_seed(args.seed)
cudnn.benchmark = True
# Redirect print to both console and log file
sys.stdout = Logger(osp.join(args.logs_dir, 'log.txt'))
# Log args
print(args)
args_dict = vars(args)
with open(osp.join(args.logs_dir, 'args.json'), 'w') as f:
json.dump(args_dict, f)
# Create data loaders
dataset, num_classes, train_loader, test_loader = \
get_data(args.dataset, args.split, args.data_dir, args.height, \
args.width, args.crop_height, args.crop_width, args.batch_size, \
args.caffe_sampler, \
args.workers)
# Create model
net_base = SSTTracker()
net_finetune = Finetune().cuda()
net_finetune.init_weights()
# Criterion
criterion = TripletLoss(margin=args.margin).cuda()
# Optimizer
optimizer = SGD_caffe(net_finetune.parameters(),
lr=args.lr,
weight_decay=args.weight_decay,
momentum=args.momentum)
# Schedule learning rate
def adjust_lr(epoch):
lr = args.lr * (0.1**(epoch // 20))
for g in optimizer.param_groups:
g['lr'] = lr
# lr = args.lr if epoch <= 200 else \
# args.lr * (0.2 ** ((epoch-200)//200 + 1))
# for g in optimizer.param_groups:
# g['lr'] = lr * g.get('lr_mult', 1)
def parse_data(inputs):
imgs, _, pids, _ = inputs
inputs = Variable(imgs).cuda()
targets = Variable(pids.cuda())
return inputs, targets
for epoch in range(750):
adjust_lr(epoch)
loss_epoch_train = 0.0
net_finetune.train()
for i, inputs in enumerate(tqdm(train_loader)):
optimizer.zero_grad()
inputs, targets = parse_data(inputs)
outputs = inputs
for k in range(35):
outputs = net_base.sst.vgg[k](outputs)
outputs = net_finetune(outputs).squeeze()
loss, prec, num_eff, num_all = criterion(outputs, targets)
loss_epoch_train += loss.data
loss.backward()
optimizer.step()
loss_epoch_train /= 41
loss_epoch_valid = 0.0
net_finetune.eval()
for i, inputs in enumerate(tqdm(test_loader)):
inputs, targets = parse_data(inputs)
outputs = inputs
for k in range(35):
outputs = net_base.sst.vgg[k](outputs)
outputs = net_finetune(outputs).squeeze()
loss, prec, num_eff, num_all = criterion(outputs, targets)
loss_epoch_valid += loss.data
loss_epoch_valid /= 108
print('epoch:%d__train_loss:%.4f__valid_loss:%.4f' %
(epoch, loss_epoch_train.data, loss_epoch_valid.data))
torch.save(net_finetune.cpu().state_dict(),
"./weights/reid_finetune_{}.pth".format(epoch + 1))
net_finetune.cuda()
def test(choice=None):
video_name_list = config['video_name_list']
if not os.path.exists(args.log_folder):
os.mkdir(args.log_folder)
save_folder = ''
choice_str = ''
if not choice is None:
choice_str = TrackerConfig.get_configure_str(choice)
save_folder = os.path.join(args.log_folder, choice_str)
if not os.path.exists(save_folder):
os.mkdir(save_folder)
save_file_name_format = os.path.join(save_folder, '{}.txt')
save_video_name_format = os.path.join(save_folder, '{}.avi')
timer = Timer()
for video_name in video_name_list:
if video_name == 'AVG-TownCentre':
TrackerConfig.set_configure((4, 0, 4, 4, 5, 4))
else:
TrackerConfig.set_configure(choice)
mot_root = os.path.join(config['mot_root'], config['dataset_type'])
mot_root = os.path.join(mot_root, video_name)
image_folder = os.path.join(mot_root, 'img1')
detection_file_name = os.path.join(mot_root, 'det/det.txt')
save_video_name = save_video_name_format.format(video_name)
save_file_name = save_file_name_format.format(video_name)
reader = MOTDataReader(image_folder=image_folder,
detection_file_name=detection_file_name,
min_confidence=None)
tracker = SSTTracker()
result = list()
result_str = save_file_name
force_init = True
for i, item in enumerate(reader):
if i > len(reader):
break
if item is None:
continue
img = item[0]
det = item[1]
if img is None or det is None or len(det) == 0:
continue
if len(det) > config['max_object']:
det = det[:config['max_object'], :]
h, w, _ = img.shape
if i == 1 and args.save_video:
vw = cv2.VideoWriter(
save_video_name,
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 10, (w, h))
det = det.astype(float)
det[:, [2, 4]] /= float(w)
det[:, [3, 5]] /= float(h)
timer.tic()
image_org = tracker.update(img, det[:, 2:6], args.show_image, i,
force_init)
force_init = False
timer.toc()
print('{}:{}, {}, {}\r'.format(os.path.basename(save_file_name), i,
int(i * 100 / len(reader)),
choice_str))
if args.show_image and not image_org is None:
cv2.imshow('res', image_org)
cv2.waitKey(1)
if args.save_video and not image_org is None:
try:
vw.write(image_org)
except:
pass
# save result
for t in tracker.tracks:
n = t.nodes[-1]
if t.age == 1:
b = n.get_box(tracker.frame_index - 1, tracker.recorder)
result.append([i + 1] + [t.id + 1] +
[b[0] * w, b[1] * h, b[2] * w, b[3] * h] +
[-1, -1, -1, -1])
# save data
np.savetxt(save_file_name, np.array(result).astype(int), fmt='%i')
print(result_str)
print(timer.total_time)
print(timer.average_time)
def test(choice=None, sequence_list=None):
image_root = args.ua_image_root
detection_root = args.ua_detection_root
ignore_root = args.ua_ignore_root
save_folder = args.save_folder
if not os.path.exists(image_root) or not os.path.exists(
detection_root) or not os.path.exists(ignore_root):
raise FileNotFoundError('Pls check the file of parameters')
print('''
==============================
= Start Reading Files =
==============================
''')
if not sequence_list is None:
sequences = np.loadtxt(sequence_list, dtype='str')
else:
sequences = os.listdir(image_root)
sequences_basename = [os.path.basename(s) for s in sequences]
# print(sequences_basename)
# validation
for seq in sequences:
if not os.path.exists(os.path.join(image_root, seq)):
raise FileNotFoundError()
all_image_folders = sorted(
[os.path.join(image_root, d) for d in sequences])
all_detection_files = [
os.path.join(detection_root,
f + '_Det_' + config['detector_name'] + '.txt')
for f in sequences_basename
]
all_ignore_files = [
os.path.join(ignore_root, f + '_IgR.txt') for f in sequences_basename
]
# all_detection_files = sorted(
# [os.path.join(detection_root, f) for f in os.listdir(detection_root) if 'MVI_' in f and os.path.basename(f) in sequences_basename]
# )
# all_ignore_files = sorted(
# [os.path.join(ignore_root, f) for f in os.listdir(ignore_root) if os.path.basename(f)[:-8] in sequences_basename]
# )
# print(all_ignore_files)
ignore_file_base_name = [
os.path.basename(f)[:-8] for f in all_ignore_files
]
detection_file_base_name = [
os.path.basename(f)[:9] for f in all_detection_files
]
choice_str = ''
if not choice is None:
choice_str = TrackerConfig.get_configure_str(choice)
TrackerConfig.set_configure(c)
save_folder = os.path.join(args.save_folder, choice_str)
if not os.path.exists(save_folder):
os.mkdir(save_folder)
saved_file_name_format = os.path.join(save_folder, '{}.txt')
saved_video_name_format = os.path.join(save_folder, '{}.avi')
for image_folder in all_image_folders:
image_folder_base_name = os.path.basename(image_folder)
i = ignore_file_base_name.index(image_folder_base_name)
ignore_file = all_ignore_files[i]
j = detection_file_base_name.index(image_folder_base_name)
detection_file = all_detection_files[j]
saved_file_name = saved_file_name_format.format(image_folder_base_name)
saved_video_name = saved_video_name_format.format(
image_folder_base_name)
vw = None
if not os.path.exists(image_folder) or not os.path.exists(
detection_file) or not os.path.exists(ignore_file):
continue
print('processing ', image_folder, '>>>>>>>>>>>>>>>>>>')
tracker = SSTTracker()
reader = UADetectionDataReader(
image_folder, detection_file,
ignore_file if args.use_ignore else None, args.detection_threshold)
result = list()
result_str = saved_file_name
timer = Timer()
for i, item in enumerate(reader):
if item is None:
continue
img = item[0]
det = item[1]
if img is None or det is None or len(det) == 0:
continue
if len(det) > config['max_object']:
det = det[:config['max_object'], :]
h, w, _ = img.shape
if vw is None and args.save_video:
vw = cv2.VideoWriter(
saved_video_name,
cv2.VideoWriter_fourcc('M', 'J', 'P', 'G'), 10, (w, h))
det[:, [2, 4]] /= float(w)
det[:, [3, 5]] /= float(h)
timer.tic()
image_org = tracker.update(img, det[:, 2:6], args.show_image, i)
timer.toc()
if i % 20 == 0:
print('{}:{}, {}, {}, {}\r'.format(
saved_file_name, i, int(i * 100 / reader.length),
choice_str, args.detection_threshold))
if args.show_image and not image_org is None:
cv2.imshow('res', image_org)
cv2.waitKey(1)
if args.save_video and not image_org is None:
vw.write(image_org)
for t in tracker.tracks:
n = t.nodes[-1]
if t.age == 1:
b = n.get_box(tracker.frame_index - 1, tracker.recorder)
result.append([i + 1] + [t.id + 1] +
[b[0] * w, b[1] * h, b[2] * w, b[3] * h] +
[-1, -1, -1, -1])
# save data
if len(result) > 0:
save_format = '%d %d %1.2f %1.2f %1.2f %1.2f %d %d %d %d'
else:
save_format = '%i'
np.savetxt(saved_file_name,
np.array(result).astype(int),
fmt=save_format)
np.savetxt(os.path.splitext(saved_file_name)[0] + '-speed.txt',
np.array([timer.total_time]),
fmt='%.3f')
print(result_str)
class DAN(Expert):
def __init__(self, model_path):
super(DAN, self).__init__("DAN")
self.model_path = model_path
def initialize(self, seq_info):
super(DAN, self).initialize(seq_info)
dataset_name = seq_info["dataset_name"]
seq_name = seq_info["seq_name"]
if dataset_name == "MOT15" and seq_name == "AVG-TownCentre":
self.choice = (4, 0, 4, 4, 5, 4)
init_test_mot15()
elif dataset_name == "MOT15":
self.choice = (0, 0, 4, 4, 5, 4)
init_test_mot15()
elif dataset_name == "DETRAC":
self.choice = TrackerConfig.get_ua_choice()
init_test_ua()
else:
self.choice = (0, 0, 4, 0, 3, 3)
init_test_mot17()
TrackerConfig.set_configure(self.choice)
if dataset_name == "MOT17":
self.min_confidence = 0.0
else:
self.min_confidence = None
config["resume"] = self.model_path
self.tracker = SSTTracker()
def track(self, img_path, dets):
super(DAN, self).track(img_path, dets)
if dets is None:
return []
img, dets, h, w = self.preprocess(img_path, dets)
if len(dets) == 0:
return []
self.tracker.update(img, dets, False, self.frame_idx)
result = []
for t in self.tracker.tracks:
n = t.nodes[-1]
if t.age == 1:
b = n.get_box(self.tracker.frame_index - 1, self.tracker.recorder)
result.append([t.id, b[0] * w, b[1] * h, b[2] * w, b[3] * h])
return result
def preprocess(self, img_path, dets):
img = cv2.imread(img_path)
if self.min_confidence is not None:
dets = dets[dets[:, 6] > self.min_confidence, :]
if len(dets) > config["max_object"]:
dets = dets[: config["max_object"], :]
h, w, _ = img.shape
dets[:, [2, 4]] /= float(w)
dets[:, [3, 5]] /= float(h)
return img, dets[:, 2:6], h, w