# -*- coding: utf-8 -*-

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

import os

import numpy as np
import tensorflow as tf
import tensorflow.contrib.eager as tfe
from tflib.utils import session
import argparse
import sys
from importlib import reload

reload(sys)


# sys.setdefaultencoding('utf8')


def batch_dataset(dataset, batch_size, prefetch_batch=2, drop_remainder=True, filter=None,
                  map_func=None, num_threads=16, shuffle=True, buffer_size=4096, repeat=-1):
    if filter:
        dataset = dataset.filter(filter)

    if map_func:
        dataset = dataset.map(map_func, num_parallel_calls=num_threads)

    if shuffle:
        dataset = dataset.shuffle(buffer_size)

    if drop_remainder:
        dataset = dataset.apply(tf.contrib.data.batch_and_drop_remainder(batch_size))
    else:
        dataset = dataset.batch(batch_size)

    dataset = dataset.repeat(repeat).prefetch(prefetch_batch)

    return dataset


def disk_image_batch_dataset(img_paths, batch_size, labels=None, prefetch_batch=2, drop_remainder=True, filter=None,
                             map_func=None, num_threads=16, shuffle=True, buffer_size=4096, repeat=-1):
    """Disk image batch dataset.

    This function is suitable for jpg and png files

    img_paths: string list or 1-D tensor, each of which is an iamge path
    labels: label list/tuple_of_list or tensor/tuple_of_tensor, each of which is a corresponding label
    """
    if labels is None:
        dataset = tf.data.Dataset.from_tensor_slices(img_paths)
    elif isinstance(labels, tuple):
        dataset = tf.data.Dataset.from_tensor_slices((img_paths,) + tuple(labels))
    else:
        dataset = tf.data.Dataset.from_tensor_slices((img_paths, labels))

    def parse_func(path, *label):
        img = tf.read_file(path)
        img = tf.image.decode_png(img, 3)
        return (img,) + label

    if map_func:
        def map_func_(*args):
            return map_func(*parse_func(*args))
    else:
        map_func_ = parse_func

    # dataset = dataset.map(parse_func, num_parallel_calls=num_threads) is slower

    dataset = batch_dataset(dataset, batch_size, prefetch_batch, drop_remainder, filter,
                            map_func_, num_threads, shuffle, buffer_size, repeat)

    return dataset


class Dataset(object):

    def __init__(self):
        self._dataset = None
        self._iterator = None
        self._batch_op = None
        self._sess = None

        self._is_eager = tf.executing_eagerly()
        self._eager_iterator = None

    def __del__(self):
        if self._sess:
            self._sess.close()

    def __iter__(self):
        return self

    def __next__(self):
        try:
            b = self.get_next()
        except:
            raise StopIteration
        else:
            return b

    next = __next__

    def get_next(self):
        if self._is_eager:
            return self._eager_iterator.get_next()
        else:
            return self._sess.run(self._batch_op)

    def reset(self, feed_dict={}):
        if self._is_eager:
            self._eager_iterator = tfe.Iterator(self._dataset)
        else:
            self._sess.run(self._iterator.initializer, feed_dict=feed_dict)

    def _bulid(self, dataset, sess=None):
        self._dataset = dataset

        if self._is_eager:
            self._eager_iterator = tfe.Iterator(dataset)
        else:
            self._iterator = dataset.make_initializable_iterator()
            self._batch_op = self._iterator.get_next()
            if sess:
                self._sess = sess
            else:
                self._sess = session()

        try:
            self.reset()
        except:
            pass

    @property
    def dataset(self):
        return self._dataset

    @property
    def iterator(self):
        return self._iterator

    @property
    def batch_op(self):
        return self._batch_op


class Celeba(Dataset):
    # att_dict = {'5_o_Clock_Shadow': 0, 'Arched_Eyebrows': 1, 'Attractive': 2,
    #             'Bags_Under_Eyes': 3, 'Bald': 4, 'Bangs': 5, 'Big_Lips': 6,
    #             'Big_Nose': 7, 'Black_Hair': 8, 'Blond_Hair': 9, 'Blurry': 10,
    #             'Brown_Hair': 11, 'Bushy_Eyebrows': 12, 'Chubby': 13,
    #             'Double_Chin': 14, 'Eyeglasses': 15, 'Goatee': 16,
    #             'Gray_Hair': 17, 'Heavy_Makeup': 18, 'High_Cheekbones': 19,
    #             'Male': 20, 'Mouth_Slightly_Open': 21, 'Mustache': 22,
    #             'Narrow_Eyes': 23, 'No_Beard': 24, 'Oval_Face': 25,
    #             'Pale_Skin': 26, 'Pointy_Nose': 27, 'Receding_Hairline': 28,
    #             'Rosy_Cheeks': 29, 'Sideburns': 30, 'Smiling': 31,
    #             'Straight_Hair': 32, 'Wavy_Hair': 33, 'Wearing_Earrings': 34,
    #             'Wearing_Hat': 35, 'Wearing_Lipstick': 36,
    #             'Wearing_Necklace': 37, 'Wearing_Necktie': 38, 'Young': 39}

    # att_dict = {
    #     "连衣裙": 0, "女式背心": 1, "女式马夹": 2, "女式针织衫": 3, "女式毛衣": 4, "女式POLO衫": 5, "女式卫衣": 6, "女式衬衫": 7, "女式T恤": 8, "女式皮草": 9,
    #     "女式皮衣": 10, "女式羽绒": 11, "女式大衣": 12, "女式风衣/长外套": 13, "女式夹克": 14, "女式西服": 15, "女式短外套": 16, "其他": 17, "泳衣泳裤": 18,
    #     "连体裤": 19, "卫衣": 20, "冲锋衣／抓绒衣": 21, "家居服/睡衣": 22, "其他": 23, "素色": 24, "拼色": 25, "渐变色": 26, "混杂色": 27,
    #     "细横条纹": 28, "粗横条纹": 29, "竖条纹": 30, "暗纹": 31, "小方格": 32, "大方格": 33, "菱形格": 34, "千鸟格 ": 35, "波点": 36, "个性印花": 37,
    #     "图案": 38, "印花": 39, "几何印花": 40, "豹纹": 41, "虎纹": 42, "斑马纹": 43, "蛇纹": 44, "迷彩": 45, "提花": 46, "其他": 47,
    #     "大理石纹": 48, "次花色模式": 49, "躯干素色": 50, "局部素色": 51, "拼色": 52, "渐变色": 53, "混杂色": 54, "细横条纹": 55, "粗横条纹": 56,
    #     "竖条纹": 57, "暗纹": 58, "小方格": 59, "大方格": 60, "菱形格": 61, "千鸟格 ": 62, "波点": 63, "个性印花": 64, "图案": 65, "印花": 66,
    #     "几何印花": 67, "豹纹": 68, "虎纹": 69, "斑马纹": 70, "蛇纹": 71, "迷彩": 72, "提花": 73, "其他": 74, "大理石纹": 75, "组合模式": 76,
    #     "拼接": 77, "叠加": 78, "图案": 79, "方块印图": 80, "字母数字": 81, "字母数字+其他": 82, "植物花卉": 83, " 人脸人像": 84, "动物卡通": 85,
    #     "骷髅头": 86, "建筑风景": 87, "其他图案": 88, "无纹理": 89, "全蕾丝": 90, "全镂空": 91, " 全网纱/透视": 92, "蕾丝拼接": 93, "镂空拼接": 94,
    #     "网纱/拼接": 95, "铰花": 96, "绒毛": 97, "压线": 98, "压线拼接": 99, "多层花边": 100, "褶皱": 101, "百褶": 102, "其他纹理": 103,
    #     "其他纹理拼接": 104, "翻领": 105, "高领": 106, "立领": 107, "毛领": 108, "西装领（平驳领/创驳领）": 109, "翻领(西装)": 110, "娃娃领/公主领": 111,
    #     "荡领": 112, "衬衫领/衬衣领": 113, "POLO领/T恤领": 114, "连帽领": 115, "棒球领": 116, "围脖领/堆领/堆堆领": 117, "系带领": 118, "青果领": 119,
    #     "圆领": 120, "V领": 121, "小V领": 122, "方领": 123, "U领": 124, "一字领": 125, "吊带领": 126, "其他无领": 127, "露肩一字领/抹胸": 128,
    #     "斜肩领": 129, "其他有领": 130, "无袖": 131, "短袖": 132, "盖肩袖": 133, "中袖/5分袖": 134, "6分袖/7分袖/8分袖": 135, "9分袖": 136,
    #     "长袖": 137, "腰部": 138, "髋部": 139, "裆部(超短裙)": 140, "大腿(短裙)": 141, "膝盖(中裙)": 142, "小腿(中长裙)": 143, "脚踝（及踝裙）": 144,
    #     "脚(及地裙)": 145, "前短后长": 146, "其它不规则": 147, "宽松": 148, "标准": 149, "修身": 150, "紧身": 151, "套头": 152, "单排扣": 153,
    #     "双排扣": 154, "居中拉链": 155, "斜拉链": 156, "暗门襟": 157, "开襟": 158, "一粒扣/暗扣": 159, "牛角扣/长条扣": 160, "裙型": 161,
    #     "鱼尾裙": 162, "蛋糕裙": 163, "节裙/塔裙": 164, "百褶裙": 165, "泡泡裙": 166, "花苞裙": 167, "蓬蓬裙": 168, "A字裙": 169, "伞裙／喇叭裙": 170,
    #     "包臀裙／一步裙": 171, "直筒裙": 172, "其他": 173, "其他": 174, "棉/麻/棉麻": 175, "雪纺": 176, "丝绸": 177, "牛仔": 178, "毛呢": 179,
    #     "毛线": 180, "皮/pu": 181, "皮草": 182, "其他": 183, "其他": 184, "红色": 185, "粉色": 186, "橙色": 187, "黄色": 188, "绿色": 189,
    #     "蓝色": 190, "紫色": 191, "灰色": 192, "黑色": 193, "白色": 194, "米色": 195, "棕色": 196, "褐色": 197, "咖色": 198, "驼色": 199,
    #     "杏色": 200, "青色": 201, "藏青色": 202, "银色": 203, "花色": 204, "其他": 205, "金色": 206, "次色": 207, "红色": 208, "粉色": 209,
    #     "橙色": 210, "黄色": 211, "绿色": 212, "蓝色": 213, "紫色": 214, "灰色": 215, "黑色": 216, "白色": 217, "米色": 218, "棕色": 219,
    #     "褐色": 220, "咖色": 221, "驼色": 222, "杏色": 223, "青色": 224, "藏青色": 225, "银色": 226, "花色": 227, "其他": 228, "金色": 229,
    #     "主元素": 230, "logo": 231, "刺绣": 232, "口袋": 233, "破洞": 234, "亮片/金属箔": 235, "铆钉/钉珠/水钻": 236, "流苏": 237, "蕾丝": 238,
    #     "镂空": 239, "补丁/贴布": 240, "系带": 241, "饰边": 242, "其他元素": 243, "镂空元素": 244, "蕾丝元素": 245, "次元素": 246, "logo": 247,
    #     "刺绣": 248, "口袋": 249, "破洞": 250, "亮片/金属箔": 251, "铆钉/钉珠/水钻": 252, "流苏": 253, "蕾丝": 254, "镂空": 255, "补丁/贴布": 256,
    #     "系带": 257, "饰边": 258, "其他元素": 259, "镂空元素": 260, "蕾丝元素": 261, "男童": 262, "男士": 263, "中性成人": 264, "女士": 265,
    #     "中性儿童": 266, "女童": 267
    # }
    #
    # att_dict = {
    #     "连衣裙": 0, "女式背心": 1, "女式马夹": 2, "女式针织衫": 3, "女式毛衣": 4, "女式POLO衫": 5, "女式卫衣": 6, "女式衬衫": 7, "女式T恤": 8, "女式皮草": 9,
    #     "女式皮衣": 10, "女式羽绒": 11, "女式大衣": 12, "女式风衣/长外套": 13, "女式夹克": 14, "女式西服": 15, "女式短外套": 16, "其他": 17, "泳衣泳裤": 18,
    #     "连体裤": 19, "卫衣": 20, "冲锋衣／抓绒衣": 21, "家居服/睡衣": 22, "其他": 23, "素色": 24, "拼色": 25, "渐变色": 26, "混杂色": 27,
    #     "细横条纹": 28, "粗横条纹": 29, "竖条纹": 30, "暗纹": 31, "小方格": 32, "大方格": 33, "菱形格": 34, "千鸟格 ": 35, "波点": 36, "个性印花": 37,
    #     "图案": 38, "印花": 39, "几何印花": 40, "豹纹": 41, "虎纹": 42, "斑马纹": 43, "蛇纹": 44, "迷彩": 45, "提花": 46, "其他": 47,
    #     "大理石纹": 48}
    #
    # att_dict = {
    #     "无袖": 131, "短袖": 132, "盖肩袖": 133, "中袖/5分袖": 134, "6分袖/7分袖/8分袖": 135, "9分袖": 136, "长袖": 137,
    #     "红色": 185, "粉色": 186, "橙色": 187, "黄色": 188, "绿色": 189,
    #     "蓝色": 190, "紫色": 191, "灰色": 192, "黑色": 193, "白色": 194, "米色": 195, "棕色": 196, "褐色": 197, "咖色": 198, "驼色": 199,
    #     "杏色": 200, "青色": 201, "藏青色": 202, "银色": 203, "花色": 204, "金色": 206
    # }

    # ----------------------------------------CHANGED-----------------------------------------

    # att_dict = {
    #     "black": 1877, "blue": 853, "brown": 281, "green": 430, "grey": 1077, "orange": 189, "pink": 435, "purple": 86,
    #     "red": 493, "white": 1595, "yellow": 173, "dresses": 0, "tops": 524, "tshirts": 2096, "shirts": 0,
    #     "blouses": 4870, "jackets": 0, "shorts": 0, "pants": 0, "jeans": 0, "skirts": 0, "jumpsuits": 0, "coats": 0
    # }
    att_dict = {
        "black": 0, "blue": 1, "brown": 2, "green": 3, "grey": 4, "orange": 5, "pink": 6, "purple": 7,
        "red": 8, "white": 9, "yellow": 10, "tops": 11, "tshirts": 12, "blouses": 13
    }

    def __init__(self, data_dir, atts, img_resize, batch_size, prefetch_batch=2, drop_remainder=True,
                 num_threads=16, shuffle=True, buffer_size=4096, repeat=-1, sess=None, part='train', crop=True):
        super(Celeba, self).__init__()

        # ----------------------------------------CHANGED-----------------------------------------
        # list_file = os.path.join(data_dir, 'list_attr_upperbody.txt')
        list_file = os.path.join(data_dir, 'AItop.txt')
        if crop:
            img_dir_jpg = os.path.join(data_dir, '')
            img_dir_png = os.path.join(data_dir, 'png')
        else:
            img_dir_jpg = os.path.join(data_dir, '')
            img_dir_png = os.path.join(data_dir, 'png')

        names = np.loadtxt(list_file, skiprows=2, usecols=[0], dtype=np.str, encoding="utf8")
        if os.path.exists(img_dir_png):
            img_paths = [os.path.join(img_dir_png, name.replace('jpg', 'png')) for name in names]
        elif os.path.exists(img_dir_jpg):
            img_paths = [os.path.join(img_dir_jpg, name) for name in names]
        else:  # not sure why add this
            img_paths = [os.path.join(img_dir_jpg, name) for name in names]

        att_id = [Celeba.att_dict[att] + 1 for att in atts]

        labels = np.loadtxt(list_file, skiprows=2, usecols=att_id, dtype=np.int64, encoding="utf8")

        if img_resize == 64:
            # crop as how VAE/GAN do
            offset_h = 40
            offset_w = 15
            img_size = 148
        else:
            offset_h = 26
            offset_w = 3
            img_size = 170

        def _map_func(img, label):
            if crop:
                img = tf.image.crop_to_bounding_box(img, offset_h, offset_w, img_size, img_size)
            # img = tf.image.resize_images(img, [img_size, img_size]) / 127.5 - 1

            # img = tf.cast(img, tf.float32)
            # img = tf.image.resize_images(img, [img_resize, img_resize]) / 127.5 - 1

            # or
            print('crop value')
            print(crop)
            img = tf.image.resize_images(img, [img_resize, img_resize], tf.image.ResizeMethod.BICUBIC)
            img = tf.clip_by_value(img, 0, 255) / 127.5 - 1
            label = (label + 1) // 2
            return img, label

        idx = np.random.permutation(len(img_paths))
        new_paths, new_labels = np.array(img_paths)[idx], np.array(labels)[idx]

        if part == 'test':
            drop_remainder = False
            shuffle = False
            repeat = 1
            img_paths = new_paths[7000:]
            labels = new_labels[7000:]
        elif part == 'val':
            img_paths = new_paths[7000:]
            labels = new_labels[7000:]
        else:
            img_paths = new_paths[:7000]
            labels = new_labels[:7000]

        dataset = disk_image_batch_dataset(img_paths=img_paths,
                                           labels=labels,
                                           batch_size=batch_size,
                                           prefetch_batch=prefetch_batch,
                                           drop_remainder=drop_remainder,
                                           map_func=_map_func,
                                           num_threads=num_threads,
                                           shuffle=shuffle,
                                           buffer_size=buffer_size,
                                           repeat=repeat)
        self._bulid(dataset, sess)

        self._img_num = len(img_paths)

    def __len__(self):
        return self._img_num

    @staticmethod
    def check_attribute_conflict(att_batch, att_name, att_names):
        def _set(att, value, att_name):
            if att_name in att_names:
                att[att_names.index(att_name)] = value

        att_id = att_names.index(att_name)

        for att in att_batch:
            # ----------------------------------------CHANGED-----------------------------------------
            # if att_name in ["无袖", "短袖", "盖肩袖", "中袖/5分袖", "6分袖/7分袖/8分袖", "9分袖", "长袖"] \
            if att_name in ["tops", "tshirts", "blouses"] \
                    and att[att_id] == 1:
                # for n in ["无袖", "短袖", "盖肩袖", "中袖/5分袖", "6分袖/7分袖/8分袖", "9分袖", "长袖"]:
                for n in ["tops", "tshirts", "blouses"]:
                    if n != att_name:
                        _set(att, 0, n)
            # elif att_name in ["红色", "粉色", "橙色", "黄色", "绿色", "蓝色", "紫色", "灰色", "黑色", "白色", "米色", "棕色", "褐色", "咖色", "驼色",
            #                   "杏色", "青色", "藏青色", "银色", "花色", "金色"] \
            elif att_name in ["black", "blue", "brown", "green", "grey", "orange", "pink", "purple", "red", "white",
                              "yellow"] \
                    and att[att_id] == 1:
                # for n in ["红色", "粉色", "橙色", "黄色", "绿色", "蓝色", "紫色", "灰色", "黑色", "白色", "米色", "棕色", "褐色", "咖色", "驼色",
                #           "杏色", "青色", "藏青色", "银色", "花色", "金色"]:
                for n in ["black", "blue", "brown", "green", "grey", "orange", "pink", "purple", "red", "white", "yellow"]:
                    if n != att_name:
                        _set(att, 0, n)

        return att_batch

    @staticmethod
    def change_attribute(att_batch, max_change_num):
        import random
        from random import randint

        for att in att_batch:
            # print('att before change')
            # print(att)
            # attr_intervals = [[0, 23], [23, 49], [49, 76], [76, 79], [79, 89], [89, 105], [105, 131], [131, 138],
            #                   [138, 148], [148, 152], [152, 161], [161, 174], [174, 184], [184, 207], [207, 230],
            #                   [230, 246], [246, 262], [262, 268]]
            attr_intervals = [[0, 10], [10, 13]]
            intervals_to_change = random.sample(attr_intervals, max_change_num)
            for intervals in intervals_to_change:
                for i in range(intervals[0], intervals[1]):
                    att[i] = 0
                change_idx = randint(intervals[0], intervals[1] - 1)
                att[change_idx] = 1

            # print('att after change')
            # print(att)
        return att_batch


if __name__ == '__main__':
    import imlib as im

    # atts = ['Bald', 'Bangs', 'Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Bushy_Eyebrows', 'Eyeglasses', 'Male',
    #         'Mouth_Slightly_Open', 'Mustache', 'No_Beard', 'Pale_Skin', 'Young']
    # atts = ["连衣裙", "女式背心", "女式马夹", "女式针织衫", "女式毛衣", "女式POLO衫", "女式卫衣", "女式衬衫", "女式T恤", "女式皮草", "女式皮衣", "女式羽绒", "女式大衣",
    #         "女式风衣/长外套", "女式夹克", "女式西服", "女式短外套", "其他", "泳衣泳裤", "连体裤", "卫衣", "冲锋衣／抓绒衣", "家居服/睡衣", "其他", "素色", "拼色",
    #         "渐变色", "混杂色", "细横条纹", "粗横条纹", "竖条纹", "暗纹", "小方格", "大方格", "菱形格", "千鸟格 ", "波点", "个性印花", "图案", "印花", "几何印花",
    #         "豹纹", "虎纹", "斑马纹", "蛇纹", "迷彩", "提花", "其他", "大理石纹", "次花色模式", "躯干素色", "局部素色", "拼色", "渐变色", "混杂色", "细横条纹",
    #         "粗横条纹", "竖条纹", "暗纹", "小方格", "大方格", "菱形格", "千鸟格 ", "波点", "个性印花", "图案", "印花", "几何印花", "豹纹", "虎纹", "斑马纹", "蛇纹",
    #         "迷彩", "提花", "其他", "大理石纹", "组合模式", "拼接", "叠加", "图案", "方块印图", "字母数字", "字母数字+其他", "植物花卉", " 人脸人像", "动物卡通",
    #         "骷髅头", "建筑风景", "其他图案", "无纹理", "全蕾丝", "全镂空", " 全网纱/透视", "蕾丝拼接", "镂空拼接", "网纱/拼接", "铰花", "绒毛", "压线", "压线拼接",
    #         "多层花边", "褶皱", "百褶", "其他纹理", "其他纹理拼接", "翻领", "高领", "立领", "毛领", "西装领（平驳领/创驳领）", "翻领(西装)", "娃娃领/公主领", "荡领",
    #         "衬衫领/衬衣领", "POLO领/T恤领", "连帽领", "棒球领", "围脖领/堆领/堆堆领", "系带领", "青果领", "圆领", "V领", "小V领", "方领", "U领", "一字领",
    #         "吊带领", "其他无领", "露肩一字领/抹胸", "斜肩领", "其他有领", "无袖", "短袖", "盖肩袖", "中袖/5分袖", "6分袖/7分袖/8分袖", "9分袖", "长袖", "腰部",
    #         "髋部", "裆部(超短裙)", "大腿(短裙)", "膝盖(中裙)", "小腿(中长裙)", "脚踝（及踝裙）", "脚(及地裙)", "前短后长", "其它不规则", "宽松", "标准", "修身",
    #         "紧身", "套头", "单排扣", "双排扣", "居中拉链", "斜拉链", "暗门襟", "开襟", "一粒扣/暗扣", "牛角扣/长条扣", "裙型", "鱼尾裙", "蛋糕裙", "节裙/塔裙",
    #         "百褶裙", "泡泡裙", "花苞裙", "蓬蓬裙", "A字裙", "伞裙／喇叭裙", "包臀裙／一步裙", "直筒裙", "其他", "其他", "棉/麻/棉麻", "雪纺", "丝绸", "牛仔", "毛呢",
    #         "毛线", "皮/pu", "皮草", "其他", "其他", "红色", "粉色", "橙色", "黄色", "绿色", "蓝色", "紫色", "灰色", "黑色", "白色", "米色", "棕色",
    #         "褐色", "咖色", "驼色", "杏色", "青色", "藏青色", "银色", "花色", "其他", "金色", "次色", "红色", "粉色", "橙色", "黄色", "绿色", "蓝色", "紫色",
    #         "灰色", "黑色", "白色", "米色", "棕色", "褐色", "咖色", "驼色", "杏色", "青色", "藏青色", "银色", "花色", "其他", "金色", "主元素", "logo",
    #         "刺绣", "口袋", "破洞", "亮片/金属箔", "铆钉/钉珠/水钻", "流苏", "蕾丝", "镂空", "补丁/贴布", "系带", "饰边", "其他元素", "镂空元素", "蕾丝元素",
    #         "次元素", "logo", "刺绣", "口袋", "破洞", "亮片/金属箔", "铆钉/钉珠/水钻", "流苏", "蕾丝", "镂空", "补丁/贴布", "系带", "饰边", "其他元素",
    #         "镂空元素", "蕾丝元素", "男童", "男士", "中性成人", "女士", "中性儿童", "女童"]
    # atts = ["无袖", "短袖", "盖肩袖", "中袖/5分袖", "6分袖/7分袖/8分袖", "9分袖", "长袖", "红色", "粉色", "橙色", "黄色", "绿色", "蓝色", "紫色", "灰色",
    #         "黑色", "白色", "米色", "棕色", "褐色", "咖色", "驼色", "杏色", "青色", "藏青色", "银色", "花色", "金色"]

    # ----------------------------------------CHANGED-----------------------------------------
    atts = ["black", "blue", "brown", "green", "grey", "orange", "pink", "purple", "red", "white", "yellow",
            "tops", "tshirts", "blouses"]

    data = Celeba('./data', atts, 128, 32, part='val')
    batch = data.get_next()
    print(len(data))
    print(batch[1][1], batch[1].dtype)
    print(batch[0].min(), batch[1].max(), batch[0].dtype)
    im.imshow(batch[0][1])
    im.show()
