def trajectory_overlap(gt_trajs, pred_traj):
"""
Calculate overlap among trajectories
"""
max_overlap = 0
max_index = -1
thresh = 0.5
for t, gt_traj in enumerate(gt_trajs):
top1 = 0
total = len(set(gt_traj.keys()) | set(pred_traj.keys()))
for i, fid in enumerate(gt_traj):
if fid not in pred_traj:
continue
sIoU = iou(gt_traj[fid], pred_traj[fid])
if sIoU >= thresh:
top1 += 1
# tIoU = (top1 + top2 + top3) * 1.0 / (3 * total)
tIoU = (top1) * 1.0 / (total)
if tIoU > max_overlap:
max_index = t
max_overlap = tIoU
return max_overlap, max_index
def merge_by_iou(preds):
"""
Just process the adjacent lines.
Args:
preds: pred_row_list
Returns:
preds: Merge some rows according to "w_iou < 0.2 and h_iou > 0.45"
"""
# col_width_max = max([len(row) for row in preds])
total_cell_count = sum([len(row) for row in preds])
i_list = [] # The index of rows merged into other rows
def take_x_min(elem):
return elem[2][0]
for i, _ in enumerate(preds):
if i == 0 or i - 1 in i_list:
continue
_, w_iou, h_iou = iou(preds[i - 1][0][2], preds[i][0][2])
if w_iou < 0.2 and h_iou > 0.45: # and len(preds[i-1]) + len(preds[i]) <= col_width_max:
preds[i - 1] += preds[i]
preds[i - 1] = sorted(preds[i - 1], key=take_x_min)
preds[i] = []
i_list.append(i)
for i in sorted(i_list,
reverse=True): # delete the row merged into other rows
if len(preds[i]) == 0:
del preds[i]
cell_count = 0 # renumber the preds
for pred in preds:
for p in pred:
cell_count += 1
p[0] = cell_count
assert total_cell_count == cell_count
# col_width_max = max([len(row) for row in preds])
# print(col_width_max)
return preds
def merge_by_iou_v2(preds):
"""
Merge rows which has only one original cell according to "w_iou < 0.2 and h_iou > 0.45"
Args:
preds: The original one cell per row
Returns:
preds: Merge some rows according to "w_iou < 0.2 and h_iou > 0.45"
"""
print(preds)
total_cell_count = sum([len(row) for row in preds])
i_list = []
def take_x_min(elem):
return elem[2][0]
for i, _ in enumerate(preds):
if i in i_list:
continue
i_tmp = i + 1
while i_tmp < total_cell_count:
print(i_tmp, total_cell_count)
_, w_iou, h_iou = iou(preds[i][0][2], preds[i_tmp][0][2])
if w_iou < 0.2 and h_iou > 0.45: # and len(preds[i-1]) + len(preds[i]) <= col_width_max:
preds[i] += preds[i_tmp]
preds[i] = sorted(preds[i], key=take_x_min)
preds[i_tmp] = []
i_list.append(i_tmp)
else:
break
i_tmp += 1
for i in sorted(i_list, reverse=True):
if len(preds[i]) == 0:
del preds[i]
print(preds)
col_width_max = max([len(row) for row in preds])
print(col_width_max)
return preds
def tubelets_overlap(tubelets_proto, annot_proto, class_idx):
for tubelet in tubelets_proto:
class_index = tubelet['class_index']
ious = []
# for each tubelet_box find the best gt_overlap
for tubelet_box in tubelet['boxes']:
tubelet_box['gt_overlap'] = 0
for annot_track in annot_proto['annotations']:
for annot_box in annot_track['track']:
if annot_box['class_index'] != class_index:
# only need to check class of first annot_box, so we break
break
if tubelet_box['frame'] == annot_box['frame']:
cur_iou = iou([annot_box['bbox']], [tubelet_box['bbox']])
# convert ndarray to a scalar
cur_iou = float(cur_iou.ravel())
if 'gt_overlap' not in tubelet_box or cur_iou > tubelet_box['gt_overlap']:
tubelet_box['gt_overlap'] = cur_iou
ious = [box['gt_overlap'] for box in tubelet['boxes']]
mean_iou = np.asarray(ious).mean()
if abs(mean_iou - 1) < np.finfo(float).eps:
tubelet['gt'] = 1
return tubelets_proto
def train_stage1(net,
optimizer,
train_loader,
val_loader,
fold=None,
weight_path='stage_1_weights',
epochs=30):
mkdir(weight_path)
log = open('log{}.txt'.format(fold), 'a+')
stat = pd.DataFrame(columns=[
'Training Loss', 'Training Acc', 'Training IoU', 'Valid Loss',
'Valid Acc', 'Valid IoU'
])
best_val_iou = .0
train_loss_record, train_acc_record, train_iou_record = [], [], []
val_loss_record, val_acc_record, val_iou_record = [], [], []
# criterion = torch.nn.BCELoss()
criterion = FocalLoss2d()
start = time.time()
print(f'\nFold-{fold} Warm-up Training Overview')
print('| Train | Val |')
print('|-----------------------------------------------------|')
print('| Loss | Acc | IoU | Loss | Acc | IoU |')
log.write(f'\nFold-{fold} Warm-up Training Overview\n')
log.write('| Train | Val |\n')
log.write('|-----------------------------------------------------|\n')
log.write('| Loss | Acc | IoU | Loss | Acc | IoU |\n')
for epoch in range(epochs):
epoch_loss = 0
epoch_acc = 0
epoch_iou = 0
for batch_idx, batch in enumerate(train_loader):
data, target = batch['image'].to(device), batch['mask'].to(device)
output = net(data)
# add sigmoid to the logits
loss = criterion(nn.Sigmoid()(output), target)
acc = accuracy(output, target)
iou_ = iou(output, target.int())
loss.backward()
optimizer.step()
optimizer.zero_grad()
epoch_loss += loss.item() / len(train_loader)
epoch_acc += acc.item() / len(train_loader)
epoch_iou += iou_.item() / len(train_loader)
print(f'|{epoch_loss:9.4f}|{epoch_acc:8.4f}|{epoch_iou:7.4f}', end='')
log.write(f'|{epoch_loss:9.4f}|{epoch_acc:8.4f}|{epoch_iou:7.4f}')
train_loss_record.append(epoch_loss)
train_acc_record.append(epoch_acc)
train_iou_record.append(epoch_iou)
with torch.no_grad():
epoch_val_loss = 0
epoch_val_acc = 0
epoch_val_iou = 0
for batch_idx, batch in enumerate(val_loader):
data, target = batch['image'].to(device), batch['mask'].to(
device)
output = net(data)
loss = criterion(nn.Sigmoid()(output), target)
acc = accuracy(output, target)
iou_ = iou(output, target.int())
epoch_val_loss += loss.item() / len(val_loader)
epoch_val_acc += acc.item() / len(val_loader)
epoch_val_iou += iou_.item() / len(val_loader)
print(
f'|{epoch_val_loss:9.4f}|{epoch_val_acc:8.4f}|{epoch_val_iou:7.4f}|'
)
log.write(
f'|{epoch_val_loss:9.4f}|{epoch_val_acc:8.4f}|{epoch_val_iou:7.4f}|\n'
)
val_loss_record.append(epoch_val_loss)
val_acc_record.append(epoch_val_acc)
val_iou_record.append(epoch_val_iou)
if epoch_val_iou > best_val_iou:
best_val_iou = epoch_val_iou
torch.save(net.state_dict(),
f'{weight_path}/Stage-1_Fold-{fold}')
stat['Epoch'] = [_ for _ in range(1, epochs + 1)]
stat['Stage'] = 1
stat['Fold'] = fold
stat['Training Loss'] = train_loss_record
stat['Training Acc'] = train_acc_record
stat['Training IoU'] = train_iou_record
stat['Valid Loss'] = val_loss_record
stat['Valid Acc'] = val_acc_record
stat['Valid IoU'] = val_iou_record
print('Time used', time.time() - start)
return net, stat
def finetune_stage(net,
train_loader,
val_loader,
fold=None,
weight_path='finetune_weights',
n_cycle=6,
initial_lr=0.01,
epochs_per_cycle=50):
mkdir(f'{weight_path}_fold{fold}')
log = open('finetune-log{}.txt'.format(fold), 'a+')
net.load_state_dict(torch.load(f'stage_2_weights/NoDepth_Fold-{fold}'))
print(f'\nModel weights from stage_2_weights/NoDepth_Fold-{fold} Loaded')
log.write(
f'\nModel weights from stage_2_weights/NoDepth_Fold-{fold} Loaded\n')
lr_record = []
train_loss_record, train_acc_record, train_iou_record = [], [], []
val_loss_record, val_acc_record, val_iou_record = [], [], []
optimizer = optim.SGD(net.parameters(),
lr=initial_lr,
momentum=0.9,
weight_decay=0.0001)
for cycle in range(n_cycle):
start = time.time()
print(f'\nFold # {fold} Snapshot # {cycle} Overview')
print('| Train | Val |')
print('|-----------------------------------------------------|')
print('| Loss | Acc | IoU | Loss | Acc | IoU |')
log.write(f'\nFold # {fold} Snapshot # {cycle} Overview\n')
log.write('| Train | Val |\n')
log.write('|-----------------------------------------------------|\n')
log.write('| Loss | Acc | IoU | Loss | Acc | IoU |\n')
best_val_iou = .0
for epoch in range(epochs_per_cycle):
epoch_loss = 0
epoch_acc = 0
epoch_iou = 0
lr = cos_annealing_lr(initial_lr, epoch, epochs_per_cycle)
if lr < 0.001:
break
lr_record.append(lr)
optimizer.state_dict()['param_groups'][0]['lr'] = lr
for batch_idx, batch in enumerate(train_loader):
data, target = batch['image'].to(device), batch['mask'].to(
device)
output = net(data)
loss = net.criterion(output, target.long())
acc = accuracy(output, target)
iou_ = iou(output, target.int())
loss.backward()
optimizer.step()
optimizer.zero_grad()
epoch_loss += loss.item() / len(train_loader)
epoch_acc += acc.item() / len(train_loader)
epoch_iou += iou_.item() / len(train_loader)
print(f'|{epoch_loss:9.4f}|{epoch_acc:8.4f}|{epoch_iou:7.4f}',
end='')
log.write(f'|{epoch_loss:9.4f}|{epoch_acc:8.4f}|{epoch_iou:7.4f}')
train_loss_record.append(epoch_loss)
train_acc_record.append(epoch_acc)
train_iou_record.append(epoch_iou)
with torch.no_grad():
epoch_val_loss = 0
epoch_val_acc = 0
epoch_val_iou = 0
for batch_idx, batch in enumerate(val_loader):
data, target = batch['image'].to(device), batch['mask'].to(
device)
output = net(data)
loss = net.criterion(output, target.long())
acc = accuracy(output, target)
iou_ = iou(output, target.int())
epoch_val_loss += loss.item() / len(val_loader)
epoch_val_acc += acc.item() / len(val_loader)
epoch_val_iou += iou_.item() / len(val_loader)
print(
f'|{epoch_val_loss:9.4f}|{epoch_val_acc:8.4f}|{epoch_val_iou:7.4f}|'
)
log.write(
f'|{epoch_val_loss:9.4f}|{epoch_val_acc:8.4f}|{epoch_val_iou:7.4f}|{epoch}\n'
)
val_loss_record.append(epoch_val_loss)
val_acc_record.append(epoch_val_acc)
val_iou_record.append(epoch_val_iou)
if epoch_val_iou > best_val_iou:
best_val_iou = epoch_val_iou
torch.save(
net.state_dict(),
f'{weight_path}_fold{fold}/cycle_{cycle}_{epoch_val_iou}'
)
print(f'Best Result: {best_val_iou}')
log.write(f'Best Result: {best_val_iou}\n')
print('Time used', time.time() - start)
def finetune_stage1(net,
train_loader,
val_loader,
fold=None,
weight_path='finetune_stage_1_weights',
epochs=80):
mkdir(weight_path)
log = open('log{}.txt'.format(fold), 'a+')
net.load_state_dict(
torch.load('stage_1_weights/Stage_1_Fold_{}'.format(fold)))
print('\nModel weights from stage_1_weights/Stage_1_Fold_{} Loaded'.format(
fold))
log.write(
'\nModel weights from stage_1_weights/Stage_1_Fold_{} Loaded\n'.format(
fold))
stat = pd.DataFrame(columns=[
'Training Loss', 'Training Acc', 'Training IoU', 'Valid Loss',
'Valid Acc', 'Valid IoU'
])
best_val_iou = .0
train_loss_record, train_acc_record, train_iou_record = [], [], []
val_loss_record, val_acc_record, val_iou_record = [], [], []
start = time.time()
optimizer = optim.SGD(net.parameters(),
lr=0.005,
momentum=0.9,
weight_decay=0.0001)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode='max',
factor=0.5,
patience=8)
print('\nFold-{} Finetune Training Overview'.format(fold))
print('| Train | Val |')
print('|-----------------------------------------------------|')
print('| Loss | Acc | IoU | Loss | Acc | IoU |')
log.write('\nFold-{} Finetune Training Overview\n'.format(fold))
log.write('| Train | Val |\n')
log.write('|-----------------------------------------------------|\n')
log.write('| Loss | Acc | IoU | Loss | Acc | IoU |\n')
early_stop_counter = 0
for epoch in range(epochs):
if early_stop_counter > 15:
print('\nEarly stop at epoch {}!'.format(epoch))
log.write('\nEarly stop at epoch {}!\n'.format(epoch))
break
epoch_loss = 0
epoch_acc = 0
epoch_iou = 0
for batch_idx, batch in enumerate(train_loader):
data, target = batch['image'].to(device), batch['mask'].to(device)
output = net(data)
# add sigmoid to the logits
loss = net.criterion(output, target)
acc = accuracy(output, target)
iou_ = iou(output, target.int())
loss.backward()
optimizer.step()
optimizer.zero_grad()
epoch_loss += loss.item() / len(train_loader)
epoch_acc += acc.item() / len(train_loader)
epoch_iou += iou_.item() / len(train_loader)
print('|{:9.4f}|{:8.4f}|{:7.4f}'.format(epoch_loss, epoch_acc,
epoch_iou),
end='')
log.write('|{:9.4f}|{:8.4f}|{:7.4f}'.format(epoch_loss, epoch_acc,
epoch_iou))
train_loss_record.append(epoch_loss)
train_acc_record.append(epoch_acc)
train_iou_record.append(epoch_iou)
with torch.no_grad():
epoch_val_loss = 0
epoch_val_acc = 0
epoch_val_iou = 0
for batch_idx, batch in enumerate(val_loader):
data, target = batch['image'].to(device), batch['mask'].to(
device)
output = net(data)
loss = net.criterion(output, target)
acc = accuracy(output, target)
iou_ = iou(output, target.int())
epoch_val_loss += loss.item() / len(val_loader)
epoch_val_acc += acc.item() / len(val_loader)
epoch_val_iou += iou_.item() / len(val_loader)
# Learning rate decay step
scheduler.step(epoch_val_iou)
print('|{:9.4f}|{:8.4f}|{:7.4f}|'.format(epoch_val_loss,
epoch_val_acc,
epoch_val_iou))
log.write('|{:9.4f}|{:8.4f}|{:7.4f}|\n'.format(
epoch_val_loss, epoch_val_acc, epoch_val_iou))
val_loss_record.append(epoch_val_loss)
val_acc_record.append(epoch_val_acc)
val_iou_record.append(epoch_val_iou)
if epoch_val_iou > best_val_iou:
early_stop_counter = 0
best_val_iou = epoch_val_iou
torch.save(net.state_dict(),
'{}/Stage_2_Fold_{}'.format(weight_path, fold))
else:
early_stop_counter += 1
# stat['Epoch'] = [_ for _ in range(1, epochs + 1)]
# stat['Stage'] = 1
# stat['Fold'] = fold
# stat['Training Loss'] = train_loss_record
# stat['Training Acc'] = train_acc_record
# stat['Training IoU'] = train_iou_record
# stat['Valid Loss'] = val_loss_record
# stat['Valid Acc'] = val_acc_record
# stat['Valid IoU'] = val_iou_record
print('Best --> {}'.format(best_val_iou))
log.write('Best --> {}\n'.format(best_val_iou))
print('Time used', time.time() - start)
return net, stat