<reponame>wwilliamcook/floop
import tensorflow as tf
from tensorflow.keras.utils import Progbar
from tensorflow.data.experimental import AUTOTUNE
import pathlib
import os
import numpy as np
import cv2 as cv
import random

from im_tools import openPDF, generate_sample, imshow


def get_paths(input_dir, glob_ext='*.*', recursive=False):
    """Return a list of the paths of all files in the given directory"""
    if recursive:
        glob_ext = '**/' + glob_ext
    data_root = pathlib.Path(input_dir)
    all_paths = list(data_root.glob(glob_ext))
    all_paths = [str(path) for path in all_paths]
    return all_paths


def generate_dataset(input_dir, output_dir, ds_size, output_res, recursive=True, show_progbar=True):
    """Take a directory of correctly-oriented PDFs and generate a dataset of rotated images and labels"""
    output_dir = pathlib.Path(output_dir)
    
    # Clear the output directory
    if output_dir.exists():
        try:
            for f in output_dir.iterdir():
                os.remove(str(f))
            output_dir.rmdir()
        except:
            print('Fatal error: unable to remove old dataset.')
            exit(-1)
    try:
        output_dir.mkdir()
    except:
        print('Fatal error: unable to create dataset directory.')
        exit(-1)

    # Get files
    all_pdf_paths = get_paths(input_dir, '*.pdf', recursive=recursive)
    print('%d PDFs found.' % len(all_pdf_paths))
    
    # Function for generating file names
    def get_write_name(index):
        return str(index).rjust(10, '0')

    # Progress bar
    if show_progbar:
        progbar = Progbar(ds_size)

    # Generate dataset
    angles = np.zeros(ds_size, dtype=np.float32)
    for i in range(ds_size):
        write_name = str(output_dir.joinpath(get_write_name(i)))
        # Loop until a sample is successfully generated
        while True:
            try:
                src_path = random.choice(all_pdf_paths)
                img, angle = generate_sample(openPDF(src_path), output_res)
                break
            except:
                pass
        # Save image and store angle
        cv.imwrite(write_name + '.jpg', img)
        angles[i] = angle
        if show_progbar:
            progbar.update(i + 1)
    if show_progbar:
        del progbar
    
    # Save angles
    with open(str(output_dir.joinpath('angles.pypickle')), 'wb+') as f:
        angles.dump(f)


def preprocess_label(label):
    label /= 360.0  # normalize to [0, 1] range

    return label


def preprocess_image(image):
    image = tf.cast(tf.image.decode_jpeg(image, channels=3), tf.float32)
    image = 2*(image / 255.0) - 1  # normalize to [-1,1] range

    return image


def load_and_preprocess_image(path):
    image = tf.io.read_file(path)
    return preprocess_image(image)


def build_dataset(input_dir, batch_size=8, training_split=.8):
    # Acquire data
    with open(str(pathlib.Path(input_dir).joinpath('angles.pypickle')), 'rb') as f:
        all_image_labels = np.load(f, allow_pickle=True)
    all_image_paths = get_paths(input_dir, '*.jpg')

    # Shuffle images
    np.random.shuffle(all_image_paths)

    # Split into training and testing datasets
    test_count = int((1 - training_split) * len(all_image_paths))
    train_count = len(all_image_paths) - test_count
    train_paths, test_paths = all_image_paths[:train_count], all_image_paths[train_count:]

    # Match the labels to the images
    def get_image_index(path):
        return int(pathlib.Path(path).name.split('.')[0])
    train_labels = [all_image_labels[get_image_index(path)] for path in train_paths]
    test_labels = [all_image_labels[get_image_index(path)] for path in test_paths]

    # Build datasets
    train_path_ds = tf.data.Dataset.from_tensor_slices(train_paths)
    test_path_ds = tf.data.Dataset.from_tensor_slices(test_paths)
    train_image_ds = train_path_ds.map(load_and_preprocess_image,
                                       num_parallel_calls=AUTOTUNE)
    test_image_ds = test_path_ds.map(load_and_preprocess_image,
                                     num_parallel_calls=AUTOTUNE)
    train_label_ds = tf.data.Dataset.from_tensor_slices(train_labels)
    test_label_ds = tf.data.Dataset.from_tensor_slices(test_labels)
#    train_label_ds = train_label_ds.map(preprocess_label, num_parallel_calls=AUTOTUNE)
#    test_label_ds = test_label_ds.map(preprocess_label, num_parallel_calls=AUTOTUNE)
    train_ds = tf.data.Dataset.zip((train_image_ds, train_label_ds))
    test_ds = tf.data.Dataset.zip((test_image_ds, test_label_ds))

    # Prepare datasets
    train_ds = train_ds.shuffle(buffer_size=train_count)
    train_ds = train_ds.repeat()
    train_ds = train_ds.batch(batch_size)
    train_ds = train_ds.prefetch(buffer_size=AUTOTUNE)

    test_ds = test_ds.shuffle(buffer_size=test_count)
    test_ds = test_ds.repeat()
    test_ds = test_ds.batch(batch_size)
    test_ds = test_ds.prefetch(buffer_size=AUTOTUNE)

    return train_ds, test_ds

import tensorflow as tf
from tensorflow.keras.utils import Progbar
from tensorflow.data.experimental import AUTOTUNE
import pathlib
import os
import numpy as np
import cv2 as cv
import random

from im_tools import openPDF, generate_sample, imshow


def get_paths(input_dir, glob_ext='*.*', recursive=False):
    """Return a list of the paths of all files in the given directory"""
    if recursive:
        glob_ext = '**/' + glob_ext
    data_root = pathlib.Path(input_dir)
    all_paths = list(data_root.glob(glob_ext))
    all_paths = [str(path) for path in all_paths]
    return all_paths


def generate_dataset(input_dir, output_dir, ds_size, output_res, recursive=True, show_progbar=True):
    """Take a directory of correctly-oriented PDFs and generate a dataset of rotated images and labels"""
    output_dir = pathlib.Path(output_dir)
    
    # Clear the output directory
    if output_dir.exists():
        try:
            for f in output_dir.iterdir():
                os.remove(str(f))
            output_dir.rmdir()
        except:
            print('Fatal error: unable to remove old dataset.')
            exit(-1)
    try:
        output_dir.mkdir()
    except:
        print('Fatal error: unable to create dataset directory.')
        exit(-1)

    # Get files
    all_pdf_paths = get_paths(input_dir, '*.pdf', recursive=recursive)
    print('%d PDFs found.' % len(all_pdf_paths))
    
    # Function for generating file names
    def get_write_name(index):
        return str(index).rjust(10, '0')

    # Progress bar
    if show_progbar:
        progbar = Progbar(ds_size)

    # Generate dataset
    angles = np.zeros(ds_size, dtype=np.float32)
    for i in range(ds_size):
        write_name = str(output_dir.joinpath(get_write_name(i)))
        # Loop until a sample is successfully generated
        while True:
            try:
                src_path = random.choice(all_pdf_paths)
                img, angle = generate_sample(openPDF(src_path), output_res)
                break
            except:
                pass
        # Save image and store angle
        cv.imwrite(write_name + '.jpg', img)
        angles[i] = angle
        if show_progbar:
            progbar.update(i + 1)
    if show_progbar:
        del progbar
    
    # Save angles
    with open(str(output_dir.joinpath('angles.pypickle')), 'wb+') as f:
        angles.dump(f)


def preprocess_label(label):
    label /= 360.0  # normalize to [0, 1] range

    return label


def preprocess_image(image):
    image = tf.cast(tf.image.decode_jpeg(image, channels=3), tf.float32)
    image = 2*(image / 255.0) - 1  # normalize to [-1,1] range

    return image


def load_and_preprocess_image(path):
    image = tf.io.read_file(path)
    return preprocess_image(image)


def build_dataset(input_dir, batch_size=8, training_split=.8):
    # Acquire data
    with open(str(pathlib.Path(input_dir).joinpath('angles.pypickle')), 'rb') as f:
        all_image_labels = np.load(f, allow_pickle=True)
    all_image_paths = get_paths(input_dir, '*.jpg')

    # Shuffle images
    np.random.shuffle(all_image_paths)

    # Split into training and testing datasets
    test_count = int((1 - training_split) * len(all_image_paths))
    train_count = len(all_image_paths) - test_count
    train_paths, test_paths = all_image_paths[:train_count], all_image_paths[train_count:]

    # Match the labels to the images
    def get_image_index(path):
        return int(pathlib.Path(path).name.split('.')[0])
    train_labels = [all_image_labels[get_image_index(path)] for path in train_paths]
    test_labels = [all_image_labels[get_image_index(path)] for path in test_paths]

    # Build datasets
    train_path_ds = tf.data.Dataset.from_tensor_slices(train_paths)
    test_path_ds = tf.data.Dataset.from_tensor_slices(test_paths)
    train_image_ds = train_path_ds.map(load_and_preprocess_image,
                                       num_parallel_calls=AUTOTUNE)
    test_image_ds = test_path_ds.map(load_and_preprocess_image,
                                     num_parallel_calls=AUTOTUNE)
    train_label_ds = tf.data.Dataset.from_tensor_slices(train_labels)
    test_label_ds = tf.data.Dataset.from_tensor_slices(test_labels)
#    train_label_ds = train_label_ds.map(preprocess_label, num_parallel_calls=AUTOTUNE)
#    test_label_ds = test_label_ds.map(preprocess_label, num_parallel_calls=AUTOTUNE)
    train_ds = tf.data.Dataset.zip((train_image_ds, train_label_ds))
    test_ds = tf.data.Dataset.zip((test_image_ds, test_label_ds))

    # Prepare datasets
    train_ds = train_ds.shuffle(buffer_size=train_count)
    train_ds = train_ds.repeat()
    train_ds = train_ds.batch(batch_size)
    train_ds = train_ds.prefetch(buffer_size=AUTOTUNE)

    test_ds = test_ds.shuffle(buffer_size=test_count)
    test_ds = test_ds.repeat()
    test_ds = test_ds.batch(batch_size)
    test_ds = test_ds.prefetch(buffer_size=AUTOTUNE)

    return train_ds, test_ds
Return a list of the paths of all files in the given directory Take a directory of correctly-oriented PDFs and generate a dataset of rotated images and labels # Clear the output directory # Get files # Function for generating file names # Progress bar # Generate dataset # Loop until a sample is successfully generated # Save image and store angle # Save angles # normalize to [0, 1] range # normalize to [-1,1] range # Acquire data # Shuffle images # Split into training and testing datasets # Match the labels to the images # Build datasets #    train_label_ds = train_label_ds.map(preprocess_label, num_parallel_calls=AUTOTUNE) #    test_label_ds = test_label_ds.map(preprocess_label, num_parallel_calls=AUTOTUNE) # Prepare datasets