from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

# import tensorflow as tf
import numpy as np
import os
import pickle
# from misc.utils import get_image
import scipy.misc
import pandas as pd

def imread(path):
    img = scipy.misc.imread(path)
    if len(img.shape) == 0:
        raise ValueError(path + " got loaded as a dimensionless array!")
    return img.astype(np.float)

def colorize(img):
    if img.ndim == 2:
        img = img.reshape(img.shape[0], img.shape[1], 1)
        img = np.concatenate([img, img, img], axis=2)
    if img.shape[2] == 4:
        img = img[:, :, 0:3]
    return img

def custom_crop(img, bbox):
    # bbox = [x-left, y-top, width, height]
    imsiz = img.shape  # [height, width, channel]
    # if box[0] + box[2] >= imsiz[1] or\
    #     box[1] + box[3] >= imsiz[0] or\
    #     box[0] <= 0 or\
    #     box[1] <= 0:
    #     box[0] = np.maximum(0, box[0])
    #     box[1] = np.maximum(0, box[1])
    #     box[2] = np.minimum(imsiz[1] - box[0] - 1, box[2])
    #     box[3] = np.minimum(imsiz[0] - box[1] - 1, box[3])
    center_x = int((2 * bbox[0] + bbox[2]) / 2)
    center_y = int((2 * bbox[1] + bbox[3]) / 2)
    R = int(np.maximum(bbox[2], bbox[3]) * 0.75)
    y1 = np.maximum(0, center_y - R)
    y2 = np.minimum(imsiz[0], center_y + R)
    x1 = np.maximum(0, center_x - R)
    x2 = np.minimum(imsiz[1], center_x + R)
    img_cropped = img[y1:y2, x1:x2, :]
    return img_cropped

def get_image(image_path, image_size, is_crop=False, bbox=None):
    global index
    out = transform(imread(image_path), image_size, is_crop, bbox)
    return out

def transform(image, image_size, is_crop, bbox):
    image = colorize(image)
    if is_crop:
        image = custom_crop(image, bbox)
    #
    transformed_image =\
        scipy.misc.imresize(image, [image_size, image_size], 'bicubic')
    return np.array(transformed_image)

# from glob import glob

# TODO: 1. current label is temporary, need to change according to real label
#       2. Current, only split the data into train, need to handel train, test

LR_HR_RETIO = 4
IMSIZE = 256
LOAD_SIZE = int(IMSIZE * 76 / 64)
BIRD_DIR = '/mnt/62645850645828D5/Ubuntu_files/causalidade_datasets/CUB_200_2011/'


def load_filenames(data_dir):
    filepath = data_dir + 'filenames.pickle'
    with open(filepath, 'rb') as f:
        filenames = pickle.load(f)
    print('Load filenames from: %s (%d)' % (filepath, len(filenames)))
    return filenames


def load_bbox(data_dir):
    bbox_path = os.path.join(data_dir, 'CUB_200_2011/bounding_boxes.txt')
    df_bounding_boxes = pd.read_csv(bbox_path,
                                    delim_whitespace=True,
                                    header=None).astype(int)
    #
    filepath = os.path.join(data_dir, 'CUB_200_2011/images.txt')
    df_filenames = pd.read_csv(filepath, delim_whitespace=True, header=None)
    filenames = df_filenames[1].tolist()
    print('Total filenames: ', len(filenames), filenames[0])
    #
    filename_bbox = {img_file[:-4]: [] for img_file in filenames}
    numImgs = len(filenames)
    for i in range(0, numImgs):
        # bbox = [x-left, y-top, width, height]
        bbox = df_bounding_boxes.iloc[i][1:].tolist()

        key = filenames[i][:-4]
        filename_bbox[key] = bbox
    #
    return filename_bbox


def save_data_list(inpath, outpath, filenames, filename_bbox):
    hr_images = []
    lr_images = []
    lr_size = int(LOAD_SIZE / LR_HR_RETIO)
    cnt = 0
    for key in filenames:
        bbox = filename_bbox[key]
        f_name = '%s/CUB_200_2011/images/%s.jpg' % (inpath, key)
        img = get_image(f_name, LOAD_SIZE, is_crop=True, bbox=bbox)
        img = img.astype('uint8')
        hr_images.append(img)
        lr_img = scipy.misc.imresize(img, [lr_size, lr_size], 'bicubic')
        lr_images.append(lr_img)

        # save image
        scipy.misc.imsave(os.path.join(outpath, "{:06d}.jpg".format(cnt+1)), img)

        cnt += 1
        if cnt % 100 == 0:
            print('Load %d......' % cnt)
    #
    print('images', len(hr_images), hr_images[0].shape, lr_images[0].shape)
    #
    # outfile = outpath + str(LOAD_SIZE) + 'images.pickle'
    # with open(outfile, 'wb') as f_out:
    #     pickle.dump(hr_images, f_out)
    #     print('save to: ', outfile)
    # #
    # outfile = outpath + str(lr_size) + 'images.pickle'
    # with open(outfile, 'wb') as f_out:
    #     pickle.dump(lr_images, f_out)
    #     print('save to: ', outfile)

def get_images_names(f):
    images = []
    for line in f:
        images.append(line.split(" ")[1].replace("\n", "").replace(".jpg", ""))
    return images


def convert_birds_dataset_pickle(inpath):
    # Load dictionary between image filename to its bbox
    filename_bbox = load_bbox(inpath)
    # print(filename_bbox)
    f_images = open(os.path.join(inpath, "CUB_200_2011", "images.txt"), "r")
    filenames = get_images_names(f_images)
    f_images.close()

    save_data_list(inpath, "/home/renato/Dev/CausalGAN/data/images/", filenames, filename_bbox)
    # ## For Train data
    # train_dir = os.path.join(inpath, 'train/')
    # train_filenames = load_filenames(train_dir)
    # save_data_list(inpath, train_dir, train_filenames, filename_bbox)

    # # ## For Test data
    # test_dir = os.path.join(inpath, 'test/')
    # test_filenames = load_filenames(test_dir)
    # save_data_list(inpath, test_dir, test_filenames, filename_bbox)


if __name__ == '__main__':
    convert_birds_dataset_pickle(BIRD_DIR)

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

# import tensorflow as tf
import numpy as np
import os
import pickle
# from misc.utils import get_image
import scipy.misc
import pandas as pd

def imread(path):
    img = scipy.misc.imread(path)
    if len(img.shape) == 0:
        raise ValueError(path + " got loaded as a dimensionless array!")
    return img.astype(np.float)

def colorize(img):
    if img.ndim == 2:
        img = img.reshape(img.shape[0], img.shape[1], 1)
        img = np.concatenate([img, img, img], axis=2)
    if img.shape[2] == 4:
        img = img[:, :, 0:3]
    return img

def custom_crop(img, bbox):
    # bbox = [x-left, y-top, width, height]
    imsiz = img.shape  # [height, width, channel]
    # if box[0] + box[2] >= imsiz[1] or\
    #     box[1] + box[3] >= imsiz[0] or\
    #     box[0] <= 0 or\
    #     box[1] <= 0:
    #     box[0] = np.maximum(0, box[0])
    #     box[1] = np.maximum(0, box[1])
    #     box[2] = np.minimum(imsiz[1] - box[0] - 1, box[2])
    #     box[3] = np.minimum(imsiz[0] - box[1] - 1, box[3])
    center_x = int((2 * bbox[0] + bbox[2]) / 2)
    center_y = int((2 * bbox[1] + bbox[3]) / 2)
    R = int(np.maximum(bbox[2], bbox[3]) * 0.75)
    y1 = np.maximum(0, center_y - R)
    y2 = np.minimum(imsiz[0], center_y + R)
    x1 = np.maximum(0, center_x - R)
    x2 = np.minimum(imsiz[1], center_x + R)
    img_cropped = img[y1:y2, x1:x2, :]
    return img_cropped

def get_image(image_path, image_size, is_crop=False, bbox=None):
    global index
    out = transform(imread(image_path), image_size, is_crop, bbox)
    return out

def transform(image, image_size, is_crop, bbox):
    image = colorize(image)
    if is_crop:
        image = custom_crop(image, bbox)
    #
    transformed_image =\
        scipy.misc.imresize(image, [image_size, image_size], 'bicubic')
    return np.array(transformed_image)

# from glob import glob

# TODO: 1. current label is temporary, need to change according to real label
#       2. Current, only split the data into train, need to handel train, test

LR_HR_RETIO = 4
IMSIZE = 256
LOAD_SIZE = int(IMSIZE * 76 / 64)
BIRD_DIR = '/mnt/62645850645828D5/Ubuntu_files/causalidade_datasets/CUB_200_2011/'


def load_filenames(data_dir):
    filepath = data_dir + 'filenames.pickle'
    with open(filepath, 'rb') as f:
        filenames = pickle.load(f)
    print('Load filenames from: %s (%d)' % (filepath, len(filenames)))
    return filenames


def load_bbox(data_dir):
    bbox_path = os.path.join(data_dir, 'CUB_200_2011/bounding_boxes.txt')
    df_bounding_boxes = pd.read_csv(bbox_path,
                                    delim_whitespace=True,
                                    header=None).astype(int)
    #
    filepath = os.path.join(data_dir, 'CUB_200_2011/images.txt')
    df_filenames = pd.read_csv(filepath, delim_whitespace=True, header=None)
    filenames = df_filenames[1].tolist()
    print('Total filenames: ', len(filenames), filenames[0])
    #
    filename_bbox = {img_file[:-4]: [] for img_file in filenames}
    numImgs = len(filenames)
    for i in range(0, numImgs):
        # bbox = [x-left, y-top, width, height]
        bbox = df_bounding_boxes.iloc[i][1:].tolist()

        key = filenames[i][:-4]
        filename_bbox[key] = bbox
    #
    return filename_bbox


def save_data_list(inpath, outpath, filenames, filename_bbox):
    hr_images = []
    lr_images = []
    lr_size = int(LOAD_SIZE / LR_HR_RETIO)
    cnt = 0
    for key in filenames:
        bbox = filename_bbox[key]
        f_name = '%s/CUB_200_2011/images/%s.jpg' % (inpath, key)
        img = get_image(f_name, LOAD_SIZE, is_crop=True, bbox=bbox)
        img = img.astype('uint8')
        hr_images.append(img)
        lr_img = scipy.misc.imresize(img, [lr_size, lr_size], 'bicubic')
        lr_images.append(lr_img)

        # save image
        scipy.misc.imsave(os.path.join(outpath, "{:06d}.jpg".format(cnt+1)), img)

        cnt += 1
        if cnt % 100 == 0:
            print('Load %d......' % cnt)
    #
    print('images', len(hr_images), hr_images[0].shape, lr_images[0].shape)
    #
    # outfile = outpath + str(LOAD_SIZE) + 'images.pickle'
    # with open(outfile, 'wb') as f_out:
    #     pickle.dump(hr_images, f_out)
    #     print('save to: ', outfile)
    # #
    # outfile = outpath + str(lr_size) + 'images.pickle'
    # with open(outfile, 'wb') as f_out:
    #     pickle.dump(lr_images, f_out)
    #     print('save to: ', outfile)

def get_images_names(f):
    images = []
    for line in f:
        images.append(line.split(" ")[1].replace("\n", "").replace(".jpg", ""))
    return images


def convert_birds_dataset_pickle(inpath):
    # Load dictionary between image filename to its bbox
    filename_bbox = load_bbox(inpath)
    # print(filename_bbox)
    f_images = open(os.path.join(inpath, "CUB_200_2011", "images.txt"), "r")
    filenames = get_images_names(f_images)
    f_images.close()

    save_data_list(inpath, "/home/renato/Dev/CausalGAN/data/images/", filenames, filename_bbox)
    # ## For Train data
    # train_dir = os.path.join(inpath, 'train/')
    # train_filenames = load_filenames(train_dir)
    # save_data_list(inpath, train_dir, train_filenames, filename_bbox)

    # # ## For Test data
    # test_dir = os.path.join(inpath, 'test/')
    # test_filenames = load_filenames(test_dir)
    # save_data_list(inpath, test_dir, test_filenames, filename_bbox)


if __name__ == '__main__':
    convert_birds_dataset_pickle(BIRD_DIR)

# import tensorflow as tf # from misc.utils import get_image # bbox = [x-left, y-top, width, height] # [height, width, channel] # if box[0] + box[2] >= imsiz[1] or\ #     box[1] + box[3] >= imsiz[0] or\ #     box[0] <= 0 or\ #     box[1] <= 0: #     box[0] = np.maximum(0, box[0]) #     box[1] = np.maximum(0, box[1]) #     box[2] = np.minimum(imsiz[1] - box[0] - 1, box[2]) #     box[3] = np.minimum(imsiz[0] - box[1] - 1, box[3]) # # from glob import glob # TODO: 1. current label is temporary, need to change according to real label #       2. Current, only split the data into train, need to handel train, test # # # bbox = [x-left, y-top, width, height] # # save image # # # outfile = outpath + str(LOAD_SIZE) + 'images.pickle' # with open(outfile, 'wb') as f_out: #     pickle.dump(hr_images, f_out) #     print('save to: ', outfile) # # # outfile = outpath + str(lr_size) + 'images.pickle' # with open(outfile, 'wb') as f_out: #     pickle.dump(lr_images, f_out) #     print('save to: ', outfile) # Load dictionary between image filename to its bbox # print(filename_bbox) # ## For Train data # train_dir = os.path.join(inpath, 'train/') # train_filenames = load_filenames(train_dir) # save_data_list(inpath, train_dir, train_filenames, filename_bbox) # # ## For Test data # test_dir = os.path.join(inpath, 'test/') # test_filenames = load_filenames(test_dir) # save_data_list(inpath, test_dir, test_filenames, filename_bbox)