Sleds/TFFirstPageEngine/FirstPageReader.py

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)