mirror of
https://github.com/leigest519/ScreenCoder.git
synced 2026-02-13 02:02:48 +00:00
73 lines
2.3 KiB
Python
73 lines
2.3 KiB
Python
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
|