def update_alignment(out, gt_lines, alignments, idx_to_char, idx_mapping,
sol_positions):
preds = out.cpu()
batch_size = preds.size(1)
preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
for i, logits in enumerate(out.data.cpu().numpy()):
raw_decode, raw_decode_full = string_utils.naive_decode(logits)
pred_str = string_utils.label2str_single(raw_decode, idx_to_char,
False)
for j, gt in enumerate(gt_lines):
cer = error_rates.cer(gt, pred_str)
global_i = idx_mapping[i]
c = sol_positions[i, 0, -1].data[0]
# alignment_error = cer
alignment_error = cer + 0.1 * (1.0 - c)
if alignment_error < alignments[j][0]:
alignments[j][0] = alignment_error
alignments[j][1] = global_i
# alignments[j][2] = out[i][:,None,:]
alignments[j][2] = None
alignments[j][3] = pred_str
def align_to_gt_lines(decoded_hw, gt_lines):
costs = []
for i in range(len(decoded_hw)):
costs.append([])
for j in range(len(gt_lines)):
pred = decoded_hw[i]
gt = gt_lines[j]
cer = error_rates.cer(gt, pred)
costs[i].append(cer)
costs = np.array(costs)
min_idx = costs.argmin(axis=0)
min_val = costs.min(axis=0)
return min_idx, min_val
def accumulate_scores(out, out_positions, xy_positions, gt_state, idx_to_char):
preds = out.transpose(0, 1).cpu()
batch_size = preds.size(1)
preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
for i, logits in enumerate(out.data.cpu().numpy()):
raw_decode, raw_decode_full = string_utils.naive_decode(logits)
pred_str = string_utils.label2str_single(raw_decode, idx_to_char,
False)
pred_str_full = string_utils.label2str_single(raw_decode_full,
idx_to_char, True)
sub_out_positions = [
o[i].data.cpu().numpy().tolist() for o in out_positions
]
sub_xy_positions = [
o[i].data.cpu().numpy().tolist() for o in xy_positions
]
for gt_obj in gt_state:
gt_text = gt_obj['gt']
cer = error_rates.cer(gt_text, pred_str)
#This is a terrible way to do this...
gt_obj['errors'] = gt_obj.get('errors', [])
gt_obj['pred'] = gt_obj.get('pred', [])
gt_obj['pred_full'] = gt_obj.get('pred_full', [])
gt_obj['path'] = gt_obj.get('path', [])
gt_obj['path_xy'] = gt_obj.get('path_xy', [])
gt_obj['errors'].append(cer)
gt_obj['pred'].append(pred_str)
gt_obj['pred_full'].append(pred_str_full)
gt_obj['path'].append(sub_out_positions)
gt_obj['path_xy'].append(sub_xy_positions)
hw.train()
for i, x in enumerate(train_dataloader):
line_imgs = Variable(x['line_imgs'].type(dtype), requires_grad=False)
labels = Variable(x['labels'], requires_grad=False)
label_lengths = Variable(x['label_lengths'], requires_grad=False)
preds = hw(line_imgs.cpu())
output_batch = preds.permute(1, 0, 2)
out = output_batch.data.cpu().numpy()
for i, gt_line in enumerate(x['gt']):
logits = out[i, ...]
pred, raw_pred = string_utils.naive_decode(logits)
pred_str = string_utils.label2str_single(pred, idx_to_char, False)
cer = error_rates.cer(gt_line, pred_str)
sum_loss += cer
steps += 1
batch_size = preds.size(1)
preds_size = Variable(torch.IntTensor([preds.size(0)] * batch_size))
loss = criterion(preds.contiguous(), labels.contiguous(), preds_size.contiguous(), label_lengths.contiguous())
# print "after"
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("Train Loss", sum_loss / steps)
print("Real Epoch", train_dataloader.epoch)
def forward_pass(x, e2e, config, thresholds, idx_to_char, update_json=False):
gt_lines = x['gt_lines']
gt = "\n".join(gt_lines)
out_original = e2e(x)
results = {}
if out_original is None:
#TODO: not a good way to handle this, but fine for now
None
gt_lines = x['gt_lines']
gt = "\n".join(gt_lines)
out_original = E2EModel.results_to_numpy(out_original)
out_original['idx'] = np.arange(out_original['sol'].shape[0])
decoded_hw, decoded_raw_hw = E2EModel.decode_handwriting(
out_original, idx_to_char)
pick, costs = E2EModel.align_to_gt_lines(decoded_hw, gt_lines)
most_ideal_pred_lines, improved_idxs = validation_utils.update_ideal_results(
pick, costs, decoded_hw, x['gt_json'])
# if update_json:
# validation_utils.save_improved_idxs(improved_idxs, decoded_hw,
# decoded_raw_hw, out_original,
# x, config[dataset_lookup]['json_folder'], config['alignment']['trim_to_sol'])
sol_thresholds = thresholds[0]
sol_thresholds_idx = range(len(sol_thresholds))
lf_nms_ranges = thresholds[1]
lf_nms_ranges_idx = range(len(lf_nms_ranges))
lf_nms_thresholds = thresholds[2]
lf_nms_thresholds_idx = range(len(lf_nms_thresholds))
most_ideal_pred_lines = "\n".join(most_ideal_pred_lines)
ideal_pred_lines = [decoded_hw[i] for i in pick]
ideal_pred_lines = "\n".join(ideal_pred_lines)
error = error_rates.cer(gt, ideal_pred_lines)
ideal_result = error
error = error_rates.cer(gt, most_ideal_pred_lines)
most_ideal_result = error
for key in itertools.product(sol_thresholds_idx, lf_nms_ranges_idx,
lf_nms_thresholds_idx):
i, j, k = key
sol_threshold = sol_thresholds[i]
lf_nms_range = lf_nms_ranges[j]
lf_nms_threshold = lf_nms_thresholds[k]
out = copy.copy(out_original)
out = E2EModel.postprocess(out,
sol_threshold=sol_threshold,
lf_nms_params={
"overlap_range": lf_nms_range,
"overlap_threshold": lf_nms_threshold
})
order = E2EModel.read_order(out)
E2EModel.filter_on_pick(out, order)
# draw_img = E2EModel.draw_output(out, img)
# cv2.imwrite("test_b_samples/test_img_{}.png".format(a), draw_img)
preds = [decoded_hw[i] for i in out['idx']]
pred = "\n".join(preds)
error = error_rates.cer(gt, pred)
results[key] = error
return results, ideal_result, most_ideal_result
def alignment_step(config, dataset_lookup=None, model_mode='best_validation', percent_range=None):
set_list = load_file_list(config['training'][dataset_lookup])
if percent_range is not None:
start = int(len(set_list) * percent_range[0])
end = int(len(set_list) * percent_range[1])
set_list = set_list[start:end]
dataset = AlignmentDataset(set_list, None)
dataloader = DataLoader(dataset, batch_size=1, shuffle=False, num_workers=0, collate_fn=alignment_dataset.collate)
char_set_path = config['network']['hw']['char_set_path']
with open(char_set_path) as f:
char_set = json.load(f)
idx_to_char = {}
for k, v in char_set['idx_to_char'].items():
idx_to_char[int(k)] = v
sol, lf, hw = init_model(config, sol_dir=model_mode, lf_dir=model_mode, hw_dir=model_mode)
e2e = E2EModel(sol, lf, hw)
dtype = torch.cuda.FloatTensor
e2e.eval()
post_processing_config = config['training']['alignment']['validation_post_processing']
sol_thresholds = post_processing_config['sol_thresholds']
sol_thresholds_idx = range(len(sol_thresholds))
lf_nms_ranges = post_processing_config['lf_nms_ranges']
lf_nms_ranges_idx = range(len(lf_nms_ranges))
lf_nms_thresholds = post_processing_config['lf_nms_thresholds']
lf_nms_thresholds_idx = range(len(lf_nms_thresholds))
results = defaultdict(list)
aligned_results = []
best_ever_results = []
prev_time = time.time()
cnt = 0
a = 0
for x in dataloader:
sys.stdout.flush()
a += 1
if a % 100 == 0:
print(a, np.mean(aligned_results))
x = x[0]
if x is None:
print("Skipping alignment because it returned None")
continue
img = x['resized_img'].numpy()[0, ...].transpose([2, 1, 0])
img = ((img + 1) * 128).astype(np.uint8)
full_img = x['full_img'].numpy()[0, ...].transpose([2, 1, 0])
full_img = ((full_img + 1) * 128).astype(np.uint8)
gt_lines = x['gt_lines']
gt = "\n".join(gt_lines)
out_original = e2e(x)
if out_original is None:
# TODO: not a good way to handle this, but fine for now
print("Possible Error: Skipping alignment on image")
continue
out_original = e2e_postprocessing.results_to_numpy(out_original)
out_original['idx'] = np.arange(out_original['sol'].shape[0])
e2e_postprocessing.trim_ends(out_original)
decoded_hw, decoded_raw_hw = e2e_postprocessing.decode_handwriting(out_original, idx_to_char)
pick, costs = e2e_postprocessing.align_to_gt_lines(decoded_hw, gt_lines)
best_ever_pred_lines, improved_idxs = validation_utils.update_ideal_results(pick, costs, decoded_hw,
x['gt_json'])
validation_utils.save_improved_idxs(improved_idxs, decoded_hw,
decoded_raw_hw, out_original,
x, config['training'][dataset_lookup]['json_folder'])
best_ever_pred_lines = "\n".join(best_ever_pred_lines)
error = error_rates.cer(gt, best_ever_pred_lines)
best_ever_results.append(error)
aligned_pred_lines = [decoded_hw[i] for i in pick]
aligned_pred_lines = "\n".join(aligned_pred_lines)
error = error_rates.cer(gt, aligned_pred_lines)
aligned_results.append(error)
if dataset_lookup == "validation_set":
# We only care about the hyperparameter postprocessing seach for the validation set
for key in itertools.product(sol_thresholds_idx, lf_nms_ranges_idx, lf_nms_thresholds_idx):
i, j, k = key
sol_threshold = sol_thresholds[i]
lf_nms_range = lf_nms_ranges[j]
lf_nms_threshold = lf_nms_thresholds[k]
out = copy.copy(out_original)
out = e2e_postprocessing.postprocess(out,
sol_threshold=sol_threshold,
lf_nms_params={
"overlap_range": lf_nms_range,
"overlap_threshold": lf_nms_threshold
})
order = e2e_postprocessing.read_order(out)
e2e_postprocessing.filter_on_pick(out, order)
e2e_postprocessing.trim_ends(out)
preds = [decoded_hw[i] for i in out['idx']]
pred = "\n".join(preds)
error = error_rates.cer(gt, pred)
results[key].append(error)
sum_results = None
if dataset_lookup == "validation_set":
# Skipping because we didn't do the hyperparameter search
sum_results = {}
for k, v in results.items():
sum_results[k] = np.mean(v)
sum_results = sorted(sum_results.items(), key=operator.itemgetter(1))
sum_results = sum_results[0]
return sum_results, np.mean(aligned_results), np.mean(best_ever_results), sol, lf, hw
def training_step(config):
char_set_path = config['network']['hw']['char_set_path']
with open(char_set_path) as f:
char_set = json.load(f)
idx_to_char = {}
for k, v in char_set['idx_to_char'].items():
idx_to_char[int(k)] = v
train_config = config['training']
allowed_training_time = train_config['lf']['reset_interval']
init_training_time = time.time()
training_set_list = load_file_list(train_config['training_set'])
train_dataset = LfDataset(training_set_list, augmentation=True)
train_dataloader = DataLoader(train_dataset,
batch_size=1,
shuffle=True,
num_workers=0,
collate_fn=lf_dataset.collate)
batches_per_epoch = int(train_config['lf']['images_per_epoch'] /
train_config['lf']['batch_size'])
train_dataloader = DatasetWrapper(train_dataloader, batches_per_epoch)
test_set_list = load_file_list(train_config['validation_set'])
test_dataset = LfDataset(
test_set_list,
random_subset_size=train_config['lf']['validation_subset_size'])
test_dataloader = DataLoader(test_dataset,
batch_size=1,
shuffle=False,
num_workers=0,
collate_fn=lf_dataset.collate)
_, lf, hw = init_model(config, only_load=['lf', 'hw'])
hw.eval()
dtype = torch.cuda.FloatTensor
lowest_loss = np.inf
lowest_loss_i = 0
for epoch in range(10000000):
lf.eval()
sum_loss = 0.0
steps = 0.0
start_time = time.time()
for step_i, x in enumerate(test_dataloader):
if x is None:
continue
#Only single batch for now
x = x[0]
if x is None:
continue
positions = [
Variable(x_i.type(dtype), requires_grad=False)[None, ...]
for x_i in x['lf_xyrs']
]
xy_positions = [
Variable(x_i.type(dtype), requires_grad=False)[None, ...]
for x_i in x['lf_xyxy']
]
img = Variable(x['img'].type(dtype), requires_grad=False)[None,
...]
#There might be a way to handle this case later,
#but for now we will skip it
if len(xy_positions) <= 1:
print("Skipping")
continue
grid_line, _, _, xy_output = lf(img,
positions[:1],
steps=len(positions),
skip_grid=False)
line = torch.nn.functional.grid_sample(img.transpose(2, 3),
grid_line)
line = line.transpose(2, 3)
predictions = hw(line)
out = predictions.permute(1, 0, 2).data.cpu().numpy()
gt_line = x['gt']
pred, raw_pred = string_utils.naive_decode(out[0])
pred_str = string_utils.label2str_single(pred, idx_to_char, False)
cer = error_rates.cer(gt_line, pred_str)
sum_loss += cer
steps += 1
# l = line[0].transpose(0,1).transpose(1,2)
# l = (l + 1)*128
# l_np = l.data.cpu().numpy()
#
# cv2.imwrite("example_line_out.png", l_np)
# print "Saved!"
# raw_input()
# loss = lf_loss.point_loss(xy_output, xy_positions)
#
# sum_loss += loss.data[0]
# steps += 1
if epoch == 0:
print("First Validation Step Complete")
print("Benchmark Validation Loss:", sum_loss / steps)
lowest_loss = sum_loss / steps
_, lf, _ = init_model(config, lf_dir='current', only_load="lf")
optimizer = torch.optim.Adam(
lf.parameters(), lr=train_config['lf']['learning_rate'])
optim_path = os.path.join(train_config['snapshot']['current'],
"lf_optim.pt")
if os.path.exists(optim_path):
print("Loading Optim Settings")
optimizer.load_state_dict(safe_load.torch_state(optim_path))
else:
print("Failed to load Optim Settings")
if lowest_loss > sum_loss / steps:
lowest_loss = sum_loss / steps
print("Saving Best")
dirname = train_config['snapshot']['best_validation']
if not len(dirname) != 0 and os.path.exists(dirname):
os.makedirs(dirname)
save_path = os.path.join(dirname, "lf.pt")
torch.save(lf.state_dict(), save_path)
lowest_loss_i = 0
test_loss = sum_loss / steps
print("Test Loss", sum_loss / steps, lowest_loss)
print("Time:", time.time() - start_time)
print("")
if allowed_training_time < (time.time() - init_training_time):
print("Out of time: Exiting...")
break
print("Epoch", epoch)
sum_loss = 0.0
steps = 0.0
lf.train()
start_time = time.time()
for x in train_dataloader:
if x is None:
continue
#Only single batch for now
x = x[0]
if x is None:
continue
positions = [
Variable(x_i.type(dtype), requires_grad=False)[None, ...]
for x_i in x['lf_xyrs']
]
xy_positions = [
Variable(x_i.type(dtype), requires_grad=False)[None, ...]
for x_i in x['lf_xyxy']
]
img = Variable(x['img'].type(dtype), requires_grad=False)[None,
...]
#There might be a way to handle this case later,
#but for now we will skip it
if len(xy_positions) <= 1:
continue
reset_interval = 4
grid_line, _, _, xy_output = lf(img,
positions[:1],
steps=len(positions),
all_positions=positions,
reset_interval=reset_interval,
randomize=True,
skip_grid=True)
loss = lf_loss.point_loss(xy_output, xy_positions)
optimizer.zero_grad()
loss.backward()
optimizer.step()
sum_loss += loss.item()
steps += 1
print("Train Loss", sum_loss / steps)
print("Real Epoch", train_dataloader.epoch)
print("Time:", time.time() - start_time)
## Save current snapshots for next iteration
print("Saving Current")
dirname = train_config['snapshot']['current']
if not len(dirname) != 0 and os.path.exists(dirname):
os.makedirs(dirname)
save_path = os.path.join(dirname, "lf.pt")
torch.save(lf.state_dict(), save_path)
optim_path = os.path.join(dirname, "lf_optim.pt")
torch.save(optimizer.state_dict(), optim_path)
def training_step(config):
hw_network_config = config['network']['hw']
train_config = config['training']
allowed_training_time = train_config['hw']['reset_interval']
init_training_time = time.time()
char_set_path = hw_network_config['char_set_path']
with open(char_set_path) as f:
char_set = json.load(f)
idx_to_char = {}
for k, v in char_set['idx_to_char'].items():
idx_to_char[int(k)] = v
training_set_list = load_file_list(train_config['training_set'])
train_dataset = HwDataset(training_set_list,
char_set['char_to_idx'],
augmentation=True,
img_height=hw_network_config['input_height'])
train_dataloader = DataLoader(train_dataset,
batch_size=train_config['hw']['batch_size'],
shuffle=False,
num_workers=0,
collate_fn=hw_dataset.collate)
batches_per_epoch = int(train_config['hw']['images_per_epoch'] /
train_config['hw']['batch_size'])
train_dataloader = DatasetWrapper(train_dataloader, batches_per_epoch)
test_set_list = load_file_list(train_config['validation_set'])
test_dataset = HwDataset(
test_set_list,
char_set['char_to_idx'],
img_height=hw_network_config['input_height'],
random_subset_size=train_config['hw']['validation_subset_size'])
test_dataloader = DataLoader(test_dataset,
batch_size=train_config['hw']['batch_size'],
shuffle=False,
num_workers=0,
collate_fn=hw_dataset.collate)
hw = cnn_lstm.create_model(hw_network_config)
hw_path = os.path.join(train_config['snapshot']['best_validation'],
"hw.pt")
hw_state = safe_load.torch_state(hw_path)
hw.load_state_dict(hw_state)
hw.cuda()
criterion = CTCLoss(zero_infinity=True)
dtype = torch.cuda.FloatTensor
lowest_loss = np.inf
lowest_loss_i = 0
for epoch in range(10000000000):
sum_loss = 0.0
steps = 0.0
hw.eval()
for x in test_dataloader:
sys.stdout.flush()
line_imgs = Variable(x['line_imgs'].type(dtype),
requires_grad=False,
volatile=True)
labels = Variable(x['labels'], requires_grad=False, volatile=True)
label_lengths = Variable(x['label_lengths'],
requires_grad=False,
volatile=True)
preds = hw(line_imgs).cpu()
output_batch = preds.permute(1, 0, 2)
out = output_batch.data.cpu().numpy()
for i, gt_line in enumerate(x['gt']):
logits = out[i, ...]
pred, raw_pred = string_utils.naive_decode(logits)
pred_str = string_utils.label2str_single(
pred, idx_to_char, False)
cer = error_rates.cer(gt_line, pred_str)
sum_loss += cer
steps += 1
if epoch == 0:
print("First Validation Step Complete")
print("Benchmark Validation CER:", sum_loss / steps)
lowest_loss = sum_loss / steps
hw = cnn_lstm.create_model(hw_network_config)
hw_path = os.path.join(train_config['snapshot']['current'],
"hw.pt")
hw_state = safe_load.torch_state(hw_path)
hw.load_state_dict(hw_state)
hw.cuda()
optimizer = torch.optim.Adam(
hw.parameters(), lr=train_config['hw']['learning_rate'])
optim_path = os.path.join(train_config['snapshot']['current'],
"hw_optim.pt")
if os.path.exists(optim_path):
print("Loading Optim Settings")
optimizer.load_state_dict(safe_load.torch_state(optim_path))
else:
print("Failed to load Optim Settings")
if lowest_loss > sum_loss / steps:
lowest_loss = sum_loss / steps
print("Saving Best")
dirname = train_config['snapshot']['best_validation']
if not len(dirname) != 0 and os.path.exists(dirname):
os.makedirs(dirname)
save_path = os.path.join(dirname, "hw.pt")
torch.save(hw.state_dict(), save_path)
lowest_loss_i = epoch
print("Test Loss", sum_loss / steps, lowest_loss)
print("")
if allowed_training_time < (time.time() - init_training_time):
print("Out of time: Exiting...")
break
print("Epoch", epoch)
sum_loss = 0.0
steps = 0.0
hw.train()
for i, x in enumerate(train_dataloader):
line_imgs = Variable(x['line_imgs'].type(dtype),
requires_grad=False)
labels = Variable(x['labels'], requires_grad=False)
label_lengths = Variable(x['label_lengths'], requires_grad=False)
preds = hw(line_imgs).cpu()
output_batch = preds.permute(1, 0, 2)
out = output_batch.data.cpu().numpy()
# if i == 0:
# for i in xrange(out.shape[0]):
# pred, pred_raw = string_utils.naive_decode(out[i,...])
# pred_str = string_utils.label2str_single(pred_raw, idx_to_char, True)
# print pred_str
for i, gt_line in enumerate(x['gt']):
logits = out[i, ...]
pred, raw_pred = string_utils.naive_decode(logits)
pred_str = string_utils.label2str_single(
pred, idx_to_char, False)
cer = error_rates.cer(gt_line, pred_str)
sum_loss += cer
steps += 1
batch_size = preds.size(1)
preds_size = Variable(torch.IntTensor([preds.size(0)] *
batch_size))
loss = criterion(preds, labels, preds_size, label_lengths)
optimizer.zero_grad()
loss.backward()
optimizer.step()
print("Train Loss", sum_loss / steps)
print("Real Epoch", train_dataloader.epoch)
## Save current snapshots for next iteration
print("Saving Current")
dirname = train_config['snapshot']['current']
if not len(dirname) != 0 and os.path.exists(dirname):
os.makedirs(dirname)
save_path = os.path.join(dirname, "hw.pt")
torch.save(hw.state_dict(), save_path)
optim_path = os.path.join(dirname, "hw_optim.pt")
torch.save(optimizer.state_dict(), optim_path)