import scipy.io as sio
from PIL import Image
import os, glob
import datetime
import shutil

running_from_path = os.getcwd()
created_images_dir = 'images'
created_labels_dir = 'labels'
data_dir = 'data'   # data_dir为脚本所在的文件夹

def hms_string(sec_elapsed):    # 格式化显示已消耗时间
    h = int(sec_elapsed / (60 * 60))
    m = int((sec_elapsed % (60 * 60)) / 60)
    s = sec_elapsed % 60.
    return "{}:{:>02}:{:>05.2f}".format(h, m, s)

def generate_dir(set_name, root_path):   # 往images和labels文件夹下生成相应的文件夹
    images_dir = os.path.join(root_path, 'images')
    annotation_dir = os.path.join(root_path, 'annotations')

    new_images_dir = os.path.join(created_images_dir, set_name)   # 将图片从原来的文件夹复制到该文件夹下
    new_annotation_dir = os.path.join(created_labels_dir, set_name)

    if not os.path.exists(new_images_dir):
        os.makedirs(new_images_dir)

    if not os.path.exists(new_annotation_dir):
        os.makedirs(new_annotation_dir)

    for img in glob.glob(os.path.join(images_dir, "*.jpg")):    # 将图片从原来的文件夹复制到新文件夹下
        shutil.copy(img, new_images_dir)

    os.chdir(annotation_dir)        # 切换到annotation的路径下
    matlab_annotations = glob.glob("*.mat")  # 仅仅包含文件名，不包含路径
    os.chdir(running_from_path)     # 切换回原来的路径

    for matfile in matlab_annotations:
        filename = matfile.split(".")[0]

        pil_image = Image.open(os.path.join(images_dir, filename+".jpg"))

        content = sio.loadmat(os.path.join(annotation_dir, matfile), matlab_compatible=False)

        boxes = content["boxes"]

        width, height = pil_image.size

        with open(os.path.join(new_annotation_dir, filename+".txt"), "w") as hs:
            for box_idx, box in enumerate(boxes.T):
                a = box[0][0][0][0]
                b = box[0][0][0][1]
                c = box[0][0][0][2]
                d = box[0][0][0][3]

                aXY = (a[0][1], a[0][0])
                bXY = (b[0][1], b[0][0])
                cXY = (c[0][1], c[0][0])
                dXY = (d[0][1], d[0][0])

                maxX = max(aXY[0], bXY[0], cXY[0], dXY[0])
                minX = min(aXY[0], bXY[0], cXY[0], dXY[0])
                maxY = max(aXY[1], bXY[1], cXY[1], dXY[1])
                minY = min(aXY[1], bXY[1], cXY[1], dXY[1])

                # clip,防止超出边界
                maxX = min(maxX, width-1)
                minX = max(minX, 0)
                maxY = min(maxY, height-1)
                minY = max(minY, 0)

                # (<absolute_x> / <image_width>)
                norm_width = (maxX - minX) / width

                # (<absolute_y> / <image_height>)
                norm_height = (maxY - minY) / height

                center_x, center_y = (maxX + minX) / 2, (maxY + minY) / 2

                norm_center_x = center_x / width
                norm_center_y = center_y / height

                if box_idx != 0:
                    hs.write("\n")

                hs.write("0 %f %f %f %f" % (norm_center_x, norm_center_y, norm_width, norm_height)) # 0表示类别

def create_txt(dirlist, filename):
    with open(filename, "w") as txtfile:   # 在data文件夹下生成txt文件
        imglist = []

        for dir in dirlist:     # dir='images/test'
            imglist.extend(glob.glob(os.path.join(dir, "*.jpg")))   # img='images/test/abc.jpg'

        for idx, img in enumerate(imglist):
            if idx != 0:
                txtfile.write("\n")
            txtfile.write(os.path.join(data_dir, img))    # 加上前缀data

if __name__ == '__main__':
    start_time = datetime.datetime.now()

    generate_dir("train", "hand_dataset/training_dataset/training_data")    # 第一个参数表示生成的文件夹的名称
    generate_dir("test", "hand_dataset/test_dataset/test_data")
    generate_dir("validation", "hand_dataset/validation_dataset/validation_data")

    create_txt((os.path.join(created_images_dir, 'train'),          # 将train和validation文件夹下的图片合并成train
                os.path.join(created_images_dir, 'validation')),
               'train.txt')
    create_txt((os.path.join(created_images_dir, 'test'), ),
               'valid.txt')

    end_time = datetime.datetime.now()
    seconds_elapsed = (end_time - start_time).total_seconds()
    print("It took {} to execute this".format(hms_string(seconds_elapsed)))