Spaces:
Running
Running
File size: 2,379 Bytes
a383d0e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
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
|