Spaces:
Running
Running
| import os | |
| import pandas as pd | |
| import json | |
| from os.path import join as pjoin | |
| import time | |
| import cv2 | |
| def save_corners(file_path, corners, compo_name, clear=True): | |
| try: | |
| df = pd.read_csv(file_path, index_col=0) | |
| except: | |
| df = pd.DataFrame(columns=['component', 'x_max', 'x_min', 'y_max', 'y_min', 'height', 'width']) | |
| if clear: | |
| df = df.drop(df.index) | |
| for corner in corners: | |
| (up_left, bottom_right) = corner | |
| c = {'component': compo_name} | |
| (c['y_min'], c['x_min']) = up_left | |
| (c['y_max'], c['x_max']) = bottom_right | |
| c['width'] = c['y_max'] - c['y_min'] | |
| c['height'] = c['x_max'] - c['x_min'] | |
| df = df.append(c, True) | |
| df.to_csv(file_path) | |
| def save_corners_json(file_path, compos): | |
| img_shape = compos[0].image_shape | |
| output = {'img_shape': img_shape, 'compos': []} | |
| f_out = open(file_path, 'w') | |
| for compo in compos: | |
| c = {'id': compo.id, 'class': compo.category} | |
| (c['column_min'], c['row_min'], c['column_max'], c['row_max']) = compo.put_bbox() | |
| c['width'] = compo.width | |
| c['height'] = compo.height | |
| output['compos'].append(c) | |
| json.dump(output, f_out, indent=4) | |
| def save_clipping(org, output_root, corners, compo_classes, compo_index): | |
| if not os.path.exists(output_root): | |
| os.mkdir(output_root) | |
| pad = 2 | |
| for i in range(len(corners)): | |
| compo = compo_classes[i] | |
| (up_left, bottom_right) = corners[i] | |
| (col_min, row_min) = up_left | |
| (col_max, row_max) = bottom_right | |
| col_min = max(col_min - pad, 0) | |
| col_max = min(col_max + pad, org.shape[1]) | |
| row_min = max(row_min - pad, 0) | |
| row_max = min(row_max + pad, org.shape[0]) | |
| # if component type already exists, index increase by 1, otherwise add this type | |
| compo_path = pjoin(output_root, compo) | |
| if compo_classes[i] not in compo_index: | |
| compo_index[compo_classes[i]] = 0 | |
| if not os.path.exists(compo_path): | |
| os.mkdir(compo_path) | |
| else: | |
| compo_index[compo_classes[i]] += 1 | |
| clip = org[row_min:row_max, col_min:col_max] | |
| cv2.imwrite(pjoin(compo_path, str(compo_index[compo_classes[i]]) + '.png'), clip) | |
| def build_directory(directory): | |
| if not os.path.exists(directory): | |
| os.mkdir(directory) | |
| return directory | |