129 lines
4.2 KiB
Python
129 lines
4.2 KiB
Python
import os
|
|
import cv2
|
|
import tensorflow as tf
|
|
import numpy as np
|
|
|
|
IMG_SIZE = 280
|
|
|
|
|
|
CLASSES = ['first_page', 'non_first_page']
|
|
|
|
|
|
def decode_one_hot(one_hot_batch):
|
|
nonzero_indicies = tf.where(tf.not_equal(
|
|
one_hot_batch, tf.zeros_like(one_hot_batch)))
|
|
reshaped_non_zero = tf.reshape(nonzero_indicies[:, -1], tf.shape(one_hot_batch)[:-1])
|
|
|
|
return reshaped_non_zero
|
|
|
|
num_first_pages = 0
|
|
|
|
num_non_first_pages = 0
|
|
|
|
|
|
# def _get_label(file_name):
|
|
# global num_first_pages, num_non_first_pages
|
|
# if '.001.png' in file_name:
|
|
# num_first_pages += 1
|
|
# return 0
|
|
# num_non_first_pages += 1
|
|
# return 1
|
|
#
|
|
#
|
|
# # def _get_img_files(train_data_root):
|
|
# # fp_imgs = []
|
|
# # non_fp_imgs = []
|
|
# #
|
|
# # for root, _, files in os.walk(train_data_root):
|
|
# # for file in files:
|
|
# # if file.endswith('.png'):
|
|
# # if '.001.png' in file:
|
|
# # fp_imgs.append(file)
|
|
# # else:
|
|
# # non_fp_imgs.append(file)
|
|
# #
|
|
# # num_fp = len(fp_imgs)
|
|
# # num_non_fp = len(non_fp_imgs)
|
|
# # shortest = num_fp if num_fp <= num_non_fp else num_non_fp
|
|
# # fp_imgs = fp_imgs[:shortest]
|
|
# # non_fp_imgs = non_fp_imgs[:shortest]
|
|
# #
|
|
# # result = fp_imgs + non_fp_imgs
|
|
# # result.sort()
|
|
# # return result
|
|
#
|
|
#
|
|
# def _read_labeled_images(train_data_root):
|
|
# # First Page ratio: 12466/46162+12466
|
|
# file_names = []
|
|
# labels = []
|
|
# max_non_fp = 12466
|
|
# cur_non_fp_cnt = 0
|
|
# for root, _, files in os.walk(train_data_root):
|
|
# for file in files:
|
|
# if file.endswith('.png'):
|
|
# is_fp = ('.001.' in file)
|
|
#
|
|
# if is_fp or cur_non_fp_cnt < max_non_fp:
|
|
# file = os.path.join(root, file)
|
|
# file_names.append(file)
|
|
# labels.append(_get_label(file))
|
|
# if not is_fp:
|
|
# cur_non_fp_cnt += 1
|
|
#
|
|
# print('%d files are first pages, %d are not' % (num_first_pages, num_non_first_pages))
|
|
#
|
|
# return file_names, labels
|
|
#
|
|
#
|
|
# def _read_images_from_disk(input_queue):
|
|
# label = input_queue[1]
|
|
# example = input_queue[0]
|
|
# return example, label
|
|
#
|
|
#
|
|
# class PageLoader:
|
|
# def __init__(self, train_data_root='/home/cdiesch/Documents/TFFirstPageClassifier/GeneratedData/', shuffle=True):
|
|
# self.train_data_root = train_data_root
|
|
#
|
|
# img_list, lbl_list = _read_labeled_images(train_data_root)
|
|
#
|
|
# test_images = []
|
|
# test_labels = []
|
|
#
|
|
# train_images = []
|
|
# train_labels = []
|
|
#
|
|
# for i in range(len(img_list)):
|
|
# if i % 3 == 0:
|
|
# train_images.append(cv2.imread(img_list[0], cv2.IMREAD_GRAYSCALE).reshape([280, 280, 1]))
|
|
# train_labels.append(lbl_list[0])
|
|
# else:
|
|
# test_images.append(cv2.imread(img_list[0], cv2.IMREAD_GRAYSCALE).reshape([280, 280, 1]))
|
|
# test_labels.append(lbl_list[0])
|
|
# del img_list[0], lbl_list[0]
|
|
#
|
|
# self.train_images = train_images
|
|
# self.train_labels = tf.one_hot(indices=train_labels, depth=2, on_value=1.0, off_value=0.0, axis=-1,
|
|
# dtype=tf.float32, name='one_hot_train_labels')
|
|
#
|
|
# # test_images = tf.convert_to_tensor(test_images, dtype=tf.float32)
|
|
# # test_labels = tf.one_hot(indices=test_labels, depth=2, on_value=1.0, off_value=0.0, axis=-1,
|
|
# # dtype=tf.float32, name='one_hot_test_labels')
|
|
#
|
|
# train_queue = tf.train.slice_input_producer([train_images, train_labels], shuffle=shuffle)
|
|
# # test_queue = tf.train.slice_input_producer([test_images, test_labels], shuffle=shuffle)
|
|
#
|
|
# self.train_image, self.train_label = _read_images_from_disk(train_queue)
|
|
# # self.test_image, self.test_label = _read_images_from_disk(test_queue)
|
|
# self.test_batch_size = 1000
|
|
#
|
|
# def get_train_batch(self, batch_size):
|
|
# return tf.train.batch([self.train_image, self.train_label], batch_size=batch_size)
|
|
#
|
|
# def _get_test_data(self):
|
|
# return tf.train.batch([self.test_image, self.test_label], batch_size=self.test_batch_size)
|
|
|
|
# test_data = property(fget=_get_test_data)
|
|
|