import argparse

import glob

import math

import tarfile

import tempfile

from datetime import datetime, timedelta

from os import remove

from os import sep as os_path_sep

from os.path import isdir, join, exists, splitext, relpath

parser = argparse.ArgumentParser(

description="This program processes imdb-wiki face dataset to produce a tfrecord file for training a YOLO style NN. " +

"The tfrecord files will have fixed size input images and yolo style output grid containing bounding boxes, gender and age")

parser.add_argument("mat_file", help="file containing metadata for imdb-wiki face data.")

parser.add_argument("data_paths", nargs='+', help="directory or tar file containing imdb-wiki face data.")

parser.add_argument("-o", "--output_path", help="directory to write output tfrecord files", required=True)

parser.add_argument("-t", "--temp_path", help="Path to base directory to use for creating temp files", default=None)

parser.add_argument("-c", "--grid_cell_size",

help="define how many pixels in input should be 1 grid cell in yolo output.",

type=int, default=32)

parser.add_argument("-s", "--input_size",

help="defines the number of pixels in each row or column you need in input image (image will have equal width and height).",

type=int, default=250)

parser.add_argument("-m", "--minimum_size",

help="minimum image width/height. Images smaller than this will be ignored.",

type=int, default=150)

parser.add_argument("--create_cropped",

help="flag to also create a dataset with cropped face images",

action="store_true")

parser.add_argument("--no_full_size",

help="flag to not create full image dataset. This only makes sense when you have flag set to create cropped faces.",

action="store_false")

parser.add_argument("-f", "--face_size",

help="defines the number of pixels in each row or column you need in Cropped Face input image (image will have equal width and height)",

type=int, default=50)

args = parser.parse_args()

mat_file = args.mat_file

data_paths = args.data_paths

REDUCE_BY = args.grid_cell_size

REQ_SIZE = args.input_size

MIN_SIZE = args.minimum_size

OUT_FILE_PATH = args.output_path

TEMP_DIR = args.temp_path

CREATE_CROPPED = args.create_cropped

CREATE_FULL = args.no_full_size

FACE_SIZE = args.face_size

for i in data_paths:

if not exists(i) or (not (i.endswith(".tar.gz") or i.endswith(".tar")) and not isdir(i)):

parser.error("Please provide a valid directory or tar file")

# importing it later when required as it takes time

print("Importing libraries and frameworks..")

import numpy as np

import tensorflow as tf

import scipy.io as sio

# from scipy.ndimage import imread

# from matplotlib.pyplot import imread

from skimage.io import imread

# from scipy.misc import imresize

from skimage.transform import resize as imresize

from skimage.color import gray2rgb

from skimage import img_as_float64

from multiprocessing import Process, Queue, cpu_count

temp_d = tempfile.TemporaryDirectory(prefix="dataset_to_tf_", suffix=".temp", dir=TEMP_DIR)

def files(tar_or_dir, ext, shuffle=False):

yield name, join(temp_d.name, tarinfo.name)

def mat_to_py_time(mat_time):

mat_time = float(mat_time)

return datetime.fromordinal(int(mat_time)) + timedelta(days=mat_time % 1) - timedelta(days=366)

def process_and_save(shared_q, index, res_q):

remove(f)

def producer(dirs, shared_q, count_q=None):

return count

if __name__ == "__main__":

# start processes for parallel processing

total_threads = min(10, cpu_count())

shared_q = Queue(total_threads)

res_q = Queue(2)

processes = []

print("Starting parallel processes..")

for i in range(total_threads):

p = Process(target=process_and_save, args=(shared_q, i, res_q))

p.start()

processes.append(p)

# this needs to be done after making processes as we don't want processes to have a copy of this large file.

m = sio.loadmat(mat_file, struct_as_record=False, squeeze_me=True)

if 'wiki' in m:

m = m['wiki']

elif 'imdb' in m:

m = m['imdb']

else:

# wait for all processes to exit

for i in range(len(processes)):

shared_q.put(None)

for p in processes:

p.join()

parser.error("Please provide a valid mat file!")

np.random.seed()

dirs = np.random.permutation(data_paths)

print("Starting producers...")

if len(dirs) > 1:

count_q = Queue(2)

producers = []

no_of_producers = min(2, len(dirs))

per_prod = len(dirs) // no_of_producers

for i in range(no_of_producers):

p = Process(target=producer, args=(dirs[i * per_prod: i * per_prod + per_prod], shared_q, count_q))

p.start()

producers.append(p)

for p in producers:

p.join()

count = 0

for i in range(len(processes)):

count += count_q.get()

else:

count = producer(dirs, shared_q)

# wait for all processes to exit

print("Waiting for processes to terminate...")

for i in range(len(processes)):

shared_q.put(None)

tcount = 0

for i in range(len(processes)):

tcount += res_q.get()

for p in processes:

p.join()

print("total files which had only 1 face:", count)

print("total_written", tcount)

temp_d.cleanup()