Spaces:
Sleeping
Sleeping
| from utils import error_rates | |
| import copy | |
| import os | |
| import cv2 | |
| import json | |
| from copy import deepcopy | |
| import numpy as np | |
| def interpolate(key1, key2, lf, lf_idx, step_percent): | |
| x0 = lf[lf_idx][key1] | |
| y0 = lf[lf_idx][key2] | |
| x1 = lf[lf_idx+1][key1] | |
| y1 = lf[lf_idx+1][key2] | |
| x = x1 * step_percent + x0 * (1.0 - step_percent) | |
| y = y1 * step_percent + y0 * (1.0 - step_percent) | |
| return x, y | |
| def get_subdivide_pt(i, pred_full, lf): | |
| percent = (float(i)+0.5) / float(len(pred_full)) | |
| lf_percent = (len(lf)-1) * percent | |
| lf_idx = int(np.floor(lf_percent)) | |
| step_percent = lf_percent - lf_idx | |
| x0, y0 = interpolate("x0", "y0", lf, lf_idx, step_percent) | |
| x1, y1 = interpolate("x1", "y1", lf, lf_idx, step_percent) | |
| return x0, y0, x1, y1 | |
| def save_improved_idxs(improved_idxs, decoded_hw, decoded_raw_hw, out, x, json_folder): | |
| output_lines = [{ | |
| "gt": gt['gt'] | |
| } for gt in x['gt_json']] | |
| # for i in improved_idxs: | |
| for i in range(len(output_lines)): | |
| if not i in improved_idxs: | |
| output_lines[i] = x['gt_json'][i] | |
| continue | |
| k = improved_idxs[i] | |
| # We want to trim the LF results | |
| # good to keep around the full length of the prediciton | |
| # so we can generate the full line-level images later | |
| # at a different resolution | |
| line_points = [] | |
| after_line_points = [] | |
| lf_path = out['lf'] | |
| end = out['ending'][k] | |
| for j in range(len(lf_path)): | |
| p = lf_path[j][k] | |
| s = out['results_scale'] | |
| if j > end: | |
| after_line_points.append({ | |
| "x0": p[0][1] * s, | |
| "x1": p[0][0] * s, | |
| "y0": p[1][1] * s, | |
| "y1": p[1][0] * s | |
| }) | |
| else: | |
| line_points.append({ | |
| "x0": p[0][1] * s, | |
| "x1": p[0][0] * s, | |
| "y0": p[1][1] * s, | |
| "y1": p[1][0] * s | |
| }) | |
| begin = out['beginning'][k] | |
| begin_f = int(np.floor(begin)) | |
| p0 = out['lf'][begin_f][k] | |
| if begin_f+1 >= len(out['lf']): | |
| p = p0 | |
| else: | |
| p1 = out['lf'][begin_f+1][k] | |
| t = begin - np.floor(begin) | |
| p = p0 * (1 - t) + p1 * t | |
| sol_point = { | |
| "x0": p[0][1] * s, | |
| "x1": p[0][0] * s, | |
| "y0": p[1][1] * s, | |
| "y1": p[1][0] * s | |
| } | |
| img_file_name = "{}_{}.png".format(x['img_key'], i) | |
| output_lines[i]['pred'] = decoded_hw[k] | |
| output_lines[i]['pred_full'] = decoded_raw_hw[k] | |
| output_lines[i]['sol'] = sol_point | |
| output_lines[i]['lf'] = line_points | |
| output_lines[i]['after_lf'] = after_line_points | |
| output_lines[i]['start_idx'] = 1 #TODO: update to backward idx | |
| output_lines[i]['hw_path'] = img_file_name | |
| line_img = out['line_imgs'][k] | |
| full_img_file_name = os.path.join(json_folder, img_file_name) | |
| cv2.imwrite(full_img_file_name, line_img) | |
| json_path = x['json_path'] | |
| with open(json_path, 'w') as f: | |
| # print('written data to:', f) | |
| json.dump(output_lines, f) | |
| def update_ideal_results(pick, costs, decoded_hw, gt_json): | |
| most_ideal_pred = [] | |
| improved_idxs = {} | |
| for i in range(len(gt_json)): | |
| gt_obj = gt_json[i] | |
| prev_pred = gt_obj.get('pred', '') | |
| gt = gt_obj['gt'] | |
| pred = decoded_hw[pick[i]] | |
| prev_cer = error_rates.cer(gt, prev_pred) | |
| cer = costs[i] | |
| if cer > prev_cer or len(pred) == 0: | |
| most_ideal_pred.append(prev_pred) | |
| continue | |
| most_ideal_pred.append(pred) | |
| improved_idxs[i] = pick[i] | |
| return most_ideal_pred, improved_idxs | |