import os
import random
import json
import numpy as np
from PIL import Image
from 口罩检测.util import train_parameters
def get_data_list(target_path,train_list_path,eval_list_path):
    #存放的类别信息
    class_detail=[]
    class_dirs = os.listdir(target_path)
    all_class_images=0 #数据集中总的图像数量
    class_label=0 #存放类别标签
    class_dim=0 #存放类别数目
    train_list=[] #训练集
    eval_list=[]  #测试集
    for class_dir in class_dirs:
        class_dim+=1
        class_detail_list={}
        eval_sum=0
        trainer_sum=0
        class_sum=0 #每个类别有多少张图片
        path = target_path+"/"+class_dir
        img_paths = os.listdir(path)
        for img_path in img_paths:
            name_path = path+"/"+img_path
            if class_sum%10==0:
                eval_sum+=1
                eval_list.append(name_path+"\t%d"%class_label+"\n")

            else:
                trainer_sum+=1
                train_list.append(name_path+"\t%d"%class_label+"\n")
            class_sum+=1
            all_class_images+=1
        class_detail_list['class_name']=class_dir
        class_detail_list['class_label']=class_label
        class_detail_list['class_eval_images']=eval_sum
        class_detail_list['class_train_images']=trainer_sum
        class_detail.append(class_detail_list)
        train_parameters['label_dict'][str(class_label)]= class_dir
        class_label+=1

    train_parameters['class_dim']=class_dim

    random.shuffle(eval_list)
    with open(eval_list_path,'a') as eval_file:
        for eval_image in eval_list:
            eval_file.write(eval_image)

    random.shuffle(train_list)
    with open(train_list_path,'a') as train_file:
        for train_item in train_list:
            train_file.write(train_item)

    #说明json的文件信息
    readJson={}
    readJson['all_class_name']=target_path
    readJson['all_class_images']= all_class_images
    readJson['class_detail']=class_detail
    jsons = json.dumps(readJson, sort_keys=True, indent=4, separators=(',', ': '))
    with open(train_parameters['readme_path'],'w') as f:
        f.write(jsons)
    print ('生成数据列表完成！')

def custom_reader(file_list):
    def reader():
        with open(file_list,'r') as f:
            lines=[line.strip() for line in f]
            for line in lines:
                img_path,label=line.strip().split('\t')
                img =Image.open(img_path)
                if img.mode!='RGB':
                    img = img.convert('RGB')
                img = img.resize((244,244),Image.BILINEAR)
                img = np.array(img).astype('float32')
                img = img.transpose((2,0,1))
                img = img/255.0
                yield  img,int(label)
    return reader

target_path=train_parameters['target_path']
train_list_path=train_parameters['train_list_path']
eval_list_path=train_parameters['eval_list_path']
batch_size=train_parameters['train_batch_size']
#每次生成数据列表前，首先清空train.txt和eval.txt
with open(train_list_path, 'w') as f:
    f.seek(0)
    f.truncate()
with open(eval_list_path, 'w') as f:
    f.seek(0)
    f.truncate()
#生成数据列表
get_data_list(target_path,train_list_path,eval_list_path)