249 lines
7.3 KiB
Python
249 lines
7.3 KiB
Python
import os
|
|
import re
|
|
import cv2
|
|
import sys
|
|
import time
|
|
import math
|
|
import datetime
|
|
import pickle
|
|
import random
|
|
import collections
|
|
import numpy as np
|
|
import ConsoleUtils
|
|
|
|
prg_name = 'PageRotation'
|
|
prg_desc = 'Randomly rotates pages to create a data-set to train on.'
|
|
prg_vers = '0.1.0'
|
|
prg_date = '2017/10/02'
|
|
prg_auth = 'Chris Diesch <cdiesch@sequencelogic.net>'
|
|
|
|
printer = ConsoleUtils.SLPrinter(prg_name)
|
|
|
|
Rotation = collections.namedtuple('Rotation', 'direction deg')
|
|
|
|
_NO_ROTATE = Rotation('Not-Rotated', 0)
|
|
_RIGHT = Rotation('Right', 270)
|
|
_UPSIDE_DOWN = Rotation('Upside-Down', 180)
|
|
_LEFT = Rotation('Left', 90)
|
|
|
|
_ROTATE_VALS = {0: _NO_ROTATE,
|
|
1: _RIGHT,
|
|
2: _UPSIDE_DOWN,
|
|
3: _LEFT}
|
|
|
|
_NUM_ROTS = len(_ROTATE_VALS)
|
|
|
|
_DEG = u'\N{DEGREE SIGN}'
|
|
|
|
_NUM_NONE = 0
|
|
_NUM_LEFT = 0
|
|
_NUM_RIGHT = 0
|
|
_NUM_UPSIDE_DOWN = 0
|
|
_NUM_TOTAL = 0
|
|
|
|
_IMG_X = 500
|
|
_IMG_Y = 500
|
|
|
|
|
|
def _get_rotation():
|
|
rotate_dir = random.randint(0, 3)
|
|
global _NUM_NONE, _NUM_RIGHT, _NUM_UPSIDE_DOWN, _NUM_LEFT, _NUM_TOTAL
|
|
if rotate_dir not in _ROTATE_VALS.keys():
|
|
raise ValueError('The value of rotate_dir is invalid (Valid range: [0, 4], Given: %d).' % rotate_dir)
|
|
|
|
res = _ROTATE_VALS[rotate_dir]
|
|
|
|
if res == _NO_ROTATE:
|
|
_NUM_NONE += 1
|
|
elif res == _RIGHT:
|
|
_NUM_RIGHT += 1
|
|
elif res == _UPSIDE_DOWN:
|
|
_NUM_UPSIDE_DOWN += 1
|
|
elif res == _LEFT:
|
|
_NUM_LEFT += 1
|
|
|
|
_NUM_TOTAL += 1
|
|
|
|
return res
|
|
|
|
|
|
def _load_images(image_root, ratio=.75, pattern=r'page[0-9]{5}\.png'):
|
|
|
|
img_files = [os.path.join(image_root, f) for f in os.listdir(image_root) if re.match(pattern, f) is not None]
|
|
random.shuffle(img_files)
|
|
print('Found %d valid image files from "%s"' % (len(img_files), image_root))
|
|
split_idx = round(len(img_files) * ratio)
|
|
|
|
train_files = img_files[:split_idx]
|
|
test_files = img_files[split_idx:]
|
|
|
|
return train_files, test_files
|
|
|
|
|
|
def _read_and_rotate(image_files):
|
|
|
|
records = {_NO_ROTATE: [],
|
|
_RIGHT: [],
|
|
_UPSIDE_DOWN: [],
|
|
_LEFT: []}
|
|
|
|
print('Processing %d images' % len(image_files))
|
|
start_time = time.time()
|
|
i = 0
|
|
for image_file in image_files:
|
|
# for j in range(4):
|
|
rand_rot = _get_rotation()
|
|
# Add a 10 degree variance to every image...
|
|
variance = random.randint(-10, 10)
|
|
# read & gray-scale image
|
|
image = cv2.imread(image_file)
|
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
# Get the size & center
|
|
height, width = image.shape[:2]
|
|
(center_w, center_h) = ((width // 2), (height // 2))
|
|
|
|
# get the rotation degree
|
|
rot_deg = rand_rot.deg + variance
|
|
# Get the rotation matrix
|
|
rot_matrix = cv2.getRotationMatrix2D((center_w, center_h), rot_deg, 1.0)
|
|
|
|
# need these to calculate the new height/width.
|
|
cos = np.abs(rot_matrix[0, 0])
|
|
sin = np.abs(rot_matrix[0, 1])
|
|
# compute new height & width (yay trig)
|
|
new_height = int((height * cos) + (width * sin))
|
|
new_width = int((height * sin) + (width * cos))
|
|
|
|
# Compute offsets
|
|
h_offset = (new_height/2) - center_h
|
|
w_offset = (new_width/2) - center_w
|
|
# Add the offsets to the rotation matrix
|
|
rot_matrix[1, 2] += h_offset
|
|
rot_matrix[0, 2] += w_offset
|
|
# Now we can do the transformation.
|
|
image = cv2.warpAffine(image, rot_matrix, (new_width, new_height))
|
|
|
|
x_scale = _IMG_X / new_width
|
|
y_scale = _IMG_Y / new_height
|
|
# Resize the new image
|
|
image = cv2.resize(image, None, fx=x_scale, fy=y_scale)
|
|
# reshape the array
|
|
image = np.reshape(image, newshape=(_IMG_Y, _IMG_X, 1))
|
|
# normalize values between 0-1.
|
|
image = image.astype(dtype='float32')
|
|
image /= float(255)
|
|
|
|
records[rand_rot].append(_create_record(image, rand_rot, rot_deg))
|
|
|
|
i += 1
|
|
if i % 25 == 0:
|
|
cur_time = time.time() - start_time
|
|
rate = float(i/cur_time)
|
|
print(' Processed %d/%d images (%.2f images/s)' % (i, len(image_files), rate))
|
|
|
|
run_time = time.time() - start_time
|
|
rate = float(len(image_files) / run_time)
|
|
|
|
rate_str = '(%.2f images/s)' % rate
|
|
|
|
if rate < 1:
|
|
rate = 1/rate
|
|
rate_str = '(.2f s/image)' % rate
|
|
|
|
print('Finished processing images in %.4f s. %s' % (run_time, rate_str))
|
|
|
|
return records
|
|
|
|
|
|
def _create_record(image, rotation, rot_deg):
|
|
|
|
height = image.shape[0]
|
|
width = image.shape[1]
|
|
|
|
rot_dir = rotation.direction
|
|
|
|
example = {'image': image,
|
|
'height': height,
|
|
'width': width,
|
|
'rotation_dir': rot_dir,
|
|
'rotation_deg': rot_deg}
|
|
|
|
return example
|
|
|
|
|
|
def _save_records(rot_recs, out_file):
|
|
|
|
for key in rot_recs:
|
|
new_name = os.path.join(out_file, '%s.pkl' % key.direction)
|
|
print('Saving records to "%s"' % new_name)
|
|
|
|
with open(new_name, 'a+b') as writer:
|
|
pickle.dump(rot_recs[key], writer)
|
|
|
|
print('Done writing record.')
|
|
|
|
|
|
def main(in_dir, out_dir):
|
|
if not os.path.exists(out_dir):
|
|
os.makedirs(out_dir)
|
|
|
|
train_file_name = os.path.join(out_dir, 'train-data')
|
|
test_file_name = os.path.join(out_dir, 'test-data')
|
|
|
|
if not os.path.exists(train_file_name):
|
|
os.mkdir(train_file_name)
|
|
if not os.path.exists(test_file_name):
|
|
os.mkdir(test_file_name)
|
|
|
|
train_files, test_files = _load_images(in_dir)
|
|
print('Loading training data')
|
|
train_records = _read_and_rotate(train_files)
|
|
del train_files
|
|
print('Loaded training data')
|
|
|
|
none_percent = float(_NUM_NONE / _NUM_TOTAL) * 100.0
|
|
right_percent = float(_NUM_RIGHT / _NUM_TOTAL) * 100.0
|
|
up_down_percent = float(_NUM_UPSIDE_DOWN / _NUM_TOTAL) * 100.0
|
|
left_percent = float(_NUM_LEFT / _NUM_TOTAL) * 100.0
|
|
|
|
print('Training data stats:')
|
|
print(' Not Rotated: %.2f%%' % none_percent)
|
|
print(' Rotated right: %.2f%%' % right_percent)
|
|
print(' Upside-down: %.2f%%' % up_down_percent)
|
|
print(' Rotated left: %.2f%%' % left_percent)
|
|
printer.write_no_prefix('')
|
|
|
|
_save_records(train_records, train_file_name)
|
|
del train_records
|
|
|
|
print('Loading training data')
|
|
test_records = _read_and_rotate(test_files)
|
|
del test_files
|
|
print('Loaded training data')
|
|
|
|
print('Saving test data to "%s"' % test_file_name)
|
|
_save_records(test_records, test_file_name)
|
|
|
|
none_percent = float(_NUM_NONE / _NUM_TOTAL) * 100.0
|
|
right_percent = float(_NUM_RIGHT / _NUM_TOTAL) * 100.0
|
|
up_down_percent = float(_NUM_UPSIDE_DOWN / _NUM_TOTAL) * 100.0
|
|
left_percent = float(_NUM_LEFT / _NUM_TOTAL) * 100.0
|
|
|
|
printer.write_no_prefix('')
|
|
printer.write_line_break()
|
|
print('Full data stats:')
|
|
print(' Not Rotated: %.2f%%' % none_percent)
|
|
print(' Rotated right: %.2f%%' % right_percent)
|
|
print(' Upside-down: %.2f%%' % up_down_percent)
|
|
print(' Rotated left: %.2f%%' % left_percent)
|
|
printer.write_no_prefix('')
|
|
|
|
|
|
if __name__ == '__main__':
|
|
print(ConsoleUtils.get_header(prg_name, prg_vers, prg_date, prg_auth))
|
|
sys.stdout = printer
|
|
input_root = r'/run/media/cdiesch/Slow_SSD/Tests/page-rotation-pdf/page-break/'
|
|
output_root = os.path.join(r'/run/media/cdiesch/Slow_SSD/Tests/tfPageRot/Input/',
|
|
datetime.datetime.now().strftime('%Y.%m.%d-%H.%M.%S'))
|
|
main(input_root, output_root)
|