#!/user/bin/python
#coding=utf-8
import os
import numpy as np
import cv2

IMAGE_SIZE = 64
# 按照指定图像大小调整尺寸
def resize_image(image, height=IMAGE_SIZE, width=IMAGE_SIZE):
    top, bottom, left, right = (0, 0, 0, 0)

    # 获取图像尺寸
    h, w, _ = image.shape

    # 对于长宽不相等的图片，找到最长的一边
    longest_edge = max(h, w)

    # 计算短边需要增加多上像素宽度使其与长边等长
    if h < longest_edge:
        dh = longest_edge - h
        top = dh // 2
        bottom = dh - top
    elif w < longest_edge:
        dw = longest_edge - w
        left = dw // 2
        right = dw - left
    else:
        pass

        # RGB颜色
    BLACK = [0, 0, 0]

    # 给图像增加边界，是图片长、宽等长，cv2.BORDER_CONSTANT指定边界颜色由value指定
    constant = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=BLACK)

    # 调整图像大小并返回
    return cv2.resize(constant, (height, width))


# 读取训练数据
images = []
labels = []
labels_name = {}


def load_dataset(path_name):
    dir_list = os.listdir(path_name)
    class_no=0
    for class_dir in dir_list:
        labels_name[class_no]=class_dir
        img_list=os.listdir(path_name+"/"+class_dir+"/")
        for img_name in img_list:
            imgpath=path_name+"/"+class_dir+"/"+img_name
            # image = cv2.imdecode(np.fromfile(imgpath, dtype=np.int), cv2.IMREAD_COLOR)
            image=cv2.imread(imgpath)
            image = resize_image(image, IMAGE_SIZE, IMAGE_SIZE)
            images.append(image)
            labels.append(class_no)
        class_no+=1
    print(labels_name)
    return images, labels, labels_name
if __name__ == '__main__':
    train_images, train_labels, labels_name = load_dataset("train_dir")
    print(train_labels)
    # print(labels_name)
    print(train_images)