import os, cv2
import hdf5storage
import numpy as np
import sys

def process_300w(root_folder, folder_name, image_name, label_name, target_size):
    image_path = os.path.join(root_folder, folder_name, image_name)
    label_path = os.path.join(root_folder, folder_name, label_name)

    with open(label_path, 'r') as ff:
        anno = ff.readlines()[3:-1]
        anno = [x.strip().split() for x in anno]
        anno = [[int(float(x[0])), int(float(x[1]))] for x in anno]
        image = cv2.imread(image_path)
        image_height, image_width, _ = image.shape
        anno_x = [x[0] for x in anno]
        anno_y = [x[1] for x in anno]
        bbox_xmin = min(anno_x)
        bbox_ymin = min(anno_y)
        bbox_xmax = max(anno_x)
        bbox_ymax = max(anno_y)
        bbox_width = bbox_xmax - bbox_xmin
        bbox_height = bbox_ymax - bbox_ymin
        scale = 1.1 
        bbox_xmin -= int((scale-1)/2*bbox_width)
        bbox_ymin -= int((scale-1)/2*bbox_height)
        bbox_width *= scale
        bbox_height *= scale
        bbox_width = int(bbox_width)
        bbox_height = int(bbox_height)
        bbox_xmin = max(bbox_xmin, 0)
        bbox_ymin = max(bbox_ymin, 0)
        bbox_width = min(bbox_width, image_width-bbox_xmin-1)
        bbox_height = min(bbox_height, image_height-bbox_ymin-1)
        anno = [[(x-bbox_xmin)/bbox_width, (y-bbox_ymin)/bbox_height] for x,y in anno]

        bbox_xmax = bbox_xmin + bbox_width
        bbox_ymax = bbox_ymin + bbox_height
        image_crop = image[bbox_ymin:bbox_ymax, bbox_xmin:bbox_xmax, :]
        image_crop = cv2.resize(image_crop, (target_size, target_size))
        return image_crop, anno
        
def process_cofw(image, bbox, anno, target_size):
    image_height, image_width, _ = image.shape
    anno_x = anno[:29]
    anno_y = anno[29:58]
    ################################
    xmin, ymin, width, height = bbox
    xmax = xmin + width -1
    ymax = ymin + height -1
    ################################
    xmin = max(xmin, 0)
    ymin = max(ymin, 0)
    xmax = min(xmax, image_width-1)
    ymax = min(ymax, image_height-1)
    anno_x = (anno_x - xmin) / (xmax - xmin)
    anno_y = (anno_y - ymin) / (ymax - ymin)
    anno = np.concatenate([anno_x.reshape(-1,1), anno_y.reshape(-1,1)], axis=1)
    anno = list(anno)
    anno = [list(x) for x in anno]
    image_crop = image[int(ymin):int(ymax), int(xmin):int(xmax), :]
    image_crop = cv2.resize(image_crop, (target_size, target_size))
    return image_crop, anno

def process_wflw(anno, target_size):
    image_name = anno[-1]
    image_path = os.path.join('..', 'data', 'WFLW', 'WFLW_images', image_name)
    image = cv2.imread(image_path)
    image_height, image_width, _ = image.shape
    lms = anno[:196]
    lms = [float(x) for x in lms]
    lms_x = lms[0::2]
    lms_y = lms[1::2]
    lms_x = [x if x >=0 else 0 for x in lms_x] 
    lms_x = [x if x <=image_width else image_width for x in lms_x] 
    lms_y = [y if y >=0 else 0 for y in lms_y] 
    lms_y = [y if y <=image_height else image_height for y in lms_y] 
    lms = [[x,y] for x,y in zip(lms_x, lms_y)]
    lms = [x for z in lms for x in z]
    bbox = anno[196:200]
    bbox = [float(x) for x in bbox]
    attrs = anno[200:206]
    attrs = np.array([int(x) for x in attrs])
    bbox_xmin, bbox_ymin, bbox_xmax, bbox_ymax = bbox

    width = bbox_xmax - bbox_xmin
    height = bbox_ymax - bbox_ymin
    scale = 1.2  
    bbox_xmin -= width * (scale-1)/2
    bbox_ymin -= height * (scale-1)/2
    bbox_xmax += width * (scale-1)/2
    bbox_ymax += height * (scale-1)/2
    bbox_xmin = max(bbox_xmin, 0)
    bbox_ymin = max(bbox_ymin, 0)
    bbox_xmax = min(bbox_xmax, image_width-1)
    bbox_ymax = min(bbox_ymax, image_height-1)
    width = bbox_xmax - bbox_xmin
    height = bbox_ymax - bbox_ymin
    image_crop = image[int(bbox_ymin):int(bbox_ymax), int(bbox_xmin):int(bbox_xmax), :]
    image_crop = cv2.resize(image_crop, (target_size, target_size))

    tmp1 = [bbox_xmin, bbox_ymin]*98
    tmp1 = np.array(tmp1)
    tmp2 = [width, height]*98
    tmp2 = np.array(tmp2)
    lms = np.array(lms) - tmp1
    lms = lms / tmp2
    lms = lms.tolist()
    lms = zip(lms[0::2], lms[1::2])
    return image_crop, list(lms) 

def process_aflw(root_folder, image_name, bbox, anno, target_size):
    image = cv2.imread(os.path.join(root_folder, 'AFLW', 'flickr', image_name))
    image_height, image_width, _ = image.shape
    anno_x = anno[:19]
    anno_y = anno[19:]
    anno_x = [x if x >=0 else 0 for x in anno_x] 
    anno_x = [x if x <=image_width else image_width for x in anno_x] 
    anno_y = [y if y >=0 else 0 for y in anno_y] 
    anno_y = [y if y <=image_height else image_height for y in anno_y] 
    anno_x_min = min(anno_x)
    anno_x_max = max(anno_x)
    anno_y_min = min(anno_y)
    anno_y_max = max(anno_y)
    xmin, xmax, ymin, ymax = bbox
    
    xmin = max(xmin, 0)
    ymin = max(ymin, 0)
    xmax = min(xmax, image_width-1)
    ymax = min(ymax, image_height-1)

    image_crop = image[int(ymin):int(ymax), int(xmin):int(xmax), :]
    image_crop = cv2.resize(image_crop, (target_size, target_size))

    anno_x = (np.array(anno_x) - xmin) / (xmax - xmin)
    anno_y = (np.array(anno_y) - ymin) / (ymax - ymin)

    anno = np.concatenate([anno_x.reshape(-1,1), anno_y.reshape(-1,1)], axis=1).flatten()
    anno = zip(anno[0::2], anno[1::2])
    return image_crop, anno

def gen_meanface(root_folder, data_name):
    with open(os.path.join(root_folder, data_name, 'train.txt'), 'r') as f:
        annos = f.readlines()
    annos = [x.strip().split()[1:] for x in annos]
    annos = [[float(x) for x in anno] for anno in annos]
    annos = np.array(annos)
    meanface = np.mean(annos, axis=0)
    meanface = meanface.tolist()
    meanface = [str(x) for x in meanface]
    
    with open(os.path.join(root_folder, data_name, 'meanface.txt'), 'w') as f:
        f.write(' '.join(meanface))

def convert_wflw(root_folder, data_name):
    with open(os.path.join('../data/WFLW/test.txt'), 'r') as f:
        annos = f.readlines()
    annos = [x.strip().split() for x in annos]
    annos_new = []
    for anno in annos:
        annos_new.append([])
        # name
        annos_new[-1].append(anno[0])
        anno = anno[1:]
        # jaw
        for i in range(17):
            annos_new[-1].append(anno[i*2*2])
            annos_new[-1].append(anno[i*2*2+1])
        # left eyebrow
        annos_new[-1].append(anno[33*2])
        annos_new[-1].append(anno[33*2+1])
        annos_new[-1].append(anno[34*2])
        annos_new[-1].append(str((float(anno[34*2+1])+float(anno[41*2+1]))/2))
        annos_new[-1].append(anno[35*2])
        annos_new[-1].append(str((float(anno[35*2+1])+float(anno[40*2+1]))/2))
        annos_new[-1].append(anno[36*2])
        annos_new[-1].append(str((float(anno[36*2+1])+float(anno[39*2+1]))/2))
        annos_new[-1].append(anno[37*2])
        annos_new[-1].append(str((float(anno[37*2+1])+float(anno[38*2+1]))/2))
        # right eyebrow
        annos_new[-1].append(anno[42*2])
        annos_new[-1].append(str((float(anno[42*2+1])+float(anno[50*2+1]))/2))
        annos_new[-1].append(anno[43*2])
        annos_new[-1].append(str((float(anno[43*2+1])+float(anno[49*2+1]))/2))
        annos_new[-1].append(anno[44*2])
        annos_new[-1].append(str((float(anno[44*2+1])+float(anno[48*2+1]))/2))
        annos_new[-1].append(anno[45*2])
        annos_new[-1].append(str((float(anno[45*2+1])+float(anno[47*2+1]))/2))
        annos_new[-1].append(anno[46*2])
        annos_new[-1].append(anno[46*2+1])
        # nose
        for i in range(51, 60):
            annos_new[-1].append(anno[i*2])
            annos_new[-1].append(anno[i*2+1])
        # left eye
        annos_new[-1].append(anno[60*2])
        annos_new[-1].append(anno[60*2+1])
        annos_new[-1].append(str(0.666*float(anno[61*2])+0.333*float(anno[62*2])))
        annos_new[-1].append(str(0.666*float(anno[61*2+1])+0.333*float(anno[62*2+1])))
        annos_new[-1].append(str(0.666*float(anno[63*2])+0.333*float(anno[62*2])))
        annos_new[-1].append(str(0.666*float(anno[63*2+1])+0.333*float(anno[62*2+1])))
        annos_new[-1].append(anno[64*2])
        annos_new[-1].append(anno[64*2+1])
        annos_new[-1].append(str(0.666*float(anno[65*2])+0.333*float(anno[66*2])))
        annos_new[-1].append(str(0.666*float(anno[65*2+1])+0.333*float(anno[66*2+1])))
        annos_new[-1].append(str(0.666*float(anno[67*2])+0.333*float(anno[66*2])))
        annos_new[-1].append(str(0.666*float(anno[67*2+1])+0.333*float(anno[66*2+1])))
        # right eye
        annos_new[-1].append(anno[68*2])
        annos_new[-1].append(anno[68*2+1])
        annos_new[-1].append(str(0.666*float(anno[69*2])+0.333*float(anno[70*2])))
        annos_new[-1].append(str(0.666*float(anno[69*2+1])+0.333*float(anno[70*2+1])))
        annos_new[-1].append(str(0.666*float(anno[71*2])+0.333*float(anno[70*2])))
        annos_new[-1].append(str(0.666*float(anno[71*2+1])+0.333*float(anno[70*2+1])))
        annos_new[-1].append(anno[72*2])
        annos_new[-1].append(anno[72*2+1])
        annos_new[-1].append(str(0.666*float(anno[73*2])+0.333*float(anno[74*2])))
        annos_new[-1].append(str(0.666*float(anno[73*2+1])+0.333*float(anno[74*2+1])))
        annos_new[-1].append(str(0.666*float(anno[75*2])+0.333*float(anno[74*2])))
        annos_new[-1].append(str(0.666*float(anno[75*2+1])+0.333*float(anno[74*2+1])))
        # mouth
        for i in range(76, 96):
            annos_new[-1].append(anno[i*2])
            annos_new[-1].append(anno[i*2+1])

    with open(os.path.join(root_folder, data_name, 'test.txt'), 'w') as f:
        for anno in annos_new:
            f.write(' '.join(anno)+'\n')


def gen_data(root_folder, data_name, target_size):
    if not os.path.exists(os.path.join(root_folder, data_name, 'images_train')):
        os.mkdir(os.path.join(root_folder, data_name, 'images_train'))
    if not os.path.exists(os.path.join(root_folder, data_name, 'images_test')):
        os.mkdir(os.path.join(root_folder, data_name, 'images_test'))

    ################################################################################################################
    if data_name == 'data_300W':
        folders_train = ['afw', 'helen/trainset', 'lfpw/trainset']
        annos_train = {}
        for folder_train in folders_train:
            all_files = sorted(os.listdir(os.path.join(root_folder, data_name, folder_train)))
            image_files = [x for x in all_files if '.pts' not in x]
            label_files = [x for x in all_files if '.pts' in x]
            assert len(image_files) == len(label_files)
            for image_name, label_name in zip(image_files, label_files):
                print(image_name)
                image_crop, anno = process_300w(os.path.join(root_folder, 'data_300W'), folder_train, image_name, label_name, target_size)
                image_crop_name = folder_train.replace('/', '_')+'_'+image_name
                cv2.imwrite(os.path.join(root_folder, data_name, 'images_train', image_crop_name), image_crop)
                annos_train[image_crop_name] = anno
        with open(os.path.join(root_folder, data_name, 'train.txt'), 'w') as f:
            for image_crop_name, anno in annos_train.items():
                f.write(image_crop_name+' ')
                for x,y in anno:
                    f.write(str(x)+' '+str(y)+' ')
                f.write('\n')
        

        folders_test = ['helen/testset', 'lfpw/testset', 'ibug']
        annos_test = {}
        for folder_test in folders_test:
            all_files = sorted(os.listdir(os.path.join(root_folder, data_name, folder_test)))
            image_files = [x for x in all_files if '.pts' not in x]
            label_files = [x for x in all_files if '.pts' in x]
            assert len(image_files) == len(label_files)
            for image_name, label_name in zip(image_files, label_files):
                print(image_name)
                image_crop, anno = process_300w(os.path.join(root_folder, 'data_300W'), folder_test, image_name, label_name, target_size)
                image_crop_name = folder_test.replace('/', '_')+'_'+image_name
                cv2.imwrite(os.path.join(root_folder, data_name, 'images_test', image_crop_name), image_crop)
                annos_test[image_crop_name] = anno
        with open(os.path.join(root_folder, data_name, 'test.txt'), 'w') as f:
            for image_crop_name, anno in annos_test.items():
                f.write(image_crop_name+' ')
                for x,y in anno:
                    f.write(str(x)+' '+str(y)+' ')
                f.write('\n')

        annos = None
        with open(os.path.join(root_folder, data_name, 'test.txt'), 'r') as f:
            annos = f.readlines()
        with open(os.path.join(root_folder, data_name, 'test_common.txt'), 'w') as f:
            for anno in annos:
                if not 'ibug' in anno:
                    f.write(anno)
        with open(os.path.join(root_folder, data_name, 'test_challenge.txt'), 'w') as f:
            for anno in annos:
                if 'ibug' in anno:
                    f.write(anno)

        gen_meanface(root_folder, data_name)
    ################################################################################################################
    elif data_name == 'COFW':
        train_file = 'COFW_train_color.mat'
        train_mat = hdf5storage.loadmat(os.path.join(root_folder, 'COFW', train_file))
        images = train_mat['IsTr']
        bboxes = train_mat['bboxesTr']
        annos = train_mat['phisTr']

        count = 1
        with open(os.path.join(root_folder, 'COFW', 'train.txt'), 'w') as f:
            for i in range(images.shape[0]):
                image = images[i, 0]
                # grayscale
                if len(image.shape) == 2:
                    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
                # swap rgb channel to bgr
                else:
                    image = image[:,:,::-1]
                bbox = bboxes[i, :]
                anno = annos[i, :]
                image_crop, anno = process_cofw(image, bbox, anno, target_size)
                pad_num = 4-len(str(count))
                image_crop_name = 'cofw_train_' + '0' * pad_num + str(count) + '.jpg'
                print(image_crop_name)
                cv2.imwrite(os.path.join(root_folder, 'COFW', 'images_train', image_crop_name), image_crop)
                f.write(image_crop_name+' ')
                for x,y in anno:
                    f.write(str(x)+' '+str(y)+' ')
                f.write('\n')
                count += 1

        test_file = 'COFW_test_color.mat'
        test_mat = hdf5storage.loadmat(os.path.join(root_folder, 'COFW', test_file))
        images = test_mat['IsT']
        bboxes = test_mat['bboxesT']
        annos = test_mat['phisT']

        count = 1
        with open(os.path.join(root_folder, 'COFW', 'test.txt'), 'w') as f:
            for i in range(images.shape[0]):
                image = images[i, 0]
                # grayscale
                if len(image.shape) == 2:
                    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
                # swap rgb channel to bgr
                else:
                    image = image[:,:,::-1]
                bbox = bboxes[i, :]
                anno = annos[i, :]
                image_crop, anno = process_cofw(image, bbox, anno, target_size)
                pad_num = 4-len(str(count))
                image_crop_name = 'cofw_test_' + '0' * pad_num + str(count) + '.jpg'
                print(image_crop_name)
                cv2.imwrite(os.path.join(root_folder, 'COFW', 'images_test', image_crop_name), image_crop)
                f.write(image_crop_name+' ')
                for x,y in anno:
                    f.write(str(x)+' '+str(y)+' ')
                f.write('\n')
                count += 1
        gen_meanface(root_folder, data_name)
    ################################################################################################################
    elif data_name == 'WFLW':
        train_file = 'list_98pt_rect_attr_train.txt'
        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_rect_attr_train_test', train_file), 'r') as f:
            annos_train = f.readlines()
        annos_train = [x.strip().split() for x in annos_train]
        count = 1
        with open(os.path.join(root_folder, 'WFLW', 'train.txt'), 'w') as f:
            for anno_train in annos_train:
                image_crop, anno = process_wflw(anno_train, target_size)
                pad_num = 4-len(str(count))
                image_crop_name = 'wflw_train_' + '0' * pad_num + str(count) + '.jpg'
                print(image_crop_name)
                cv2.imwrite(os.path.join(root_folder, 'WFLW', 'images_train', image_crop_name), image_crop)
                f.write(image_crop_name+' ')
                for x,y in anno:
                    f.write(str(x)+' '+str(y)+' ')
                f.write('\n')
                count += 1

        test_file = 'list_98pt_rect_attr_test.txt'
        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_rect_attr_train_test', test_file), 'r') as f:
            annos_test = f.readlines()
        annos_test = [x.strip().split() for x in annos_test]
        names_mapping = {}
        count = 1
        with open(os.path.join(root_folder, 'WFLW', 'test.txt'), 'w') as f:
            for anno_test in annos_test:
                image_crop, anno = process_wflw(anno_test, target_size)
                pad_num = 4-len(str(count))
                image_crop_name = 'wflw_test_' + '0' * pad_num + str(count) + '.jpg'
                print(image_crop_name)
                names_mapping[anno_test[0]+'_'+anno_test[-1]] = [image_crop_name, anno]
                cv2.imwrite(os.path.join(root_folder, 'WFLW', 'images_test', image_crop_name), image_crop)
                f.write(image_crop_name+' ')
                for x,y in list(anno):
                    f.write(str(x)+' '+str(y)+' ')
                f.write('\n')
                count += 1

        test_pose_file = 'list_98pt_test_largepose.txt'
        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_pose_file), 'r') as f:
            annos_pose_test = f.readlines()
        names_pose = [x.strip().split() for x in annos_pose_test]
        names_pose = [x[0]+'_'+x[-1] for x in names_pose]
        with open(os.path.join(root_folder, 'WFLW', 'test_pose.txt'), 'w') as f:
            for name_pose in names_pose:
                if name_pose in names_mapping:
                    image_crop_name, anno = names_mapping[name_pose]
                    f.write(image_crop_name+' ')
                    for x,y in anno:
                        f.write(str(x)+' '+str(y)+' ')
                    f.write('\n')
                else:
                    print('error!')
                    exit(0)

        test_expr_file = 'list_98pt_test_expression.txt'
        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_expr_file), 'r') as f:
            annos_expr_test = f.readlines()
        names_expr = [x.strip().split() for x in annos_expr_test]
        names_expr = [x[0]+'_'+x[-1] for x in names_expr]
        with open(os.path.join(root_folder, 'WFLW', 'test_expr.txt'), 'w') as f:
            for name_expr in names_expr:
                if name_expr in names_mapping:
                    image_crop_name, anno = names_mapping[name_expr]
                    f.write(image_crop_name+' ')
                    for x,y in anno:
                        f.write(str(x)+' '+str(y)+' ')
                    f.write('\n')
                else:
                    print('error!')
                    exit(0)

        test_illu_file = 'list_98pt_test_illumination.txt'
        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_illu_file), 'r') as f:
            annos_illu_test = f.readlines()
        names_illu = [x.strip().split() for x in annos_illu_test]
        names_illu = [x[0]+'_'+x[-1] for x in names_illu]
        with open(os.path.join(root_folder, 'WFLW', 'test_illu.txt'), 'w') as f:
            for name_illu in names_illu:
                if name_illu in names_mapping:
                    image_crop_name, anno = names_mapping[name_illu]
                    f.write(image_crop_name+' ')
                    for x,y in anno:
                        f.write(str(x)+' '+str(y)+' ')
                    f.write('\n')
                else:
                    print('error!')
                    exit(0)

        test_mu_file = 'list_98pt_test_makeup.txt'
        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_mu_file), 'r') as f:
            annos_mu_test = f.readlines()
        names_mu = [x.strip().split() for x in annos_mu_test]
        names_mu = [x[0]+'_'+x[-1] for x in names_mu]
        with open(os.path.join(root_folder, 'WFLW', 'test_mu.txt'), 'w') as f:
            for name_mu in names_mu:
                if name_mu in names_mapping:
                    image_crop_name, anno = names_mapping[name_mu]
                    f.write(image_crop_name+' ')
                    for x,y in anno:
                        f.write(str(x)+' '+str(y)+' ')
                    f.write('\n')
                else:
                    print('error!')
                    exit(0)

        test_occu_file = 'list_98pt_test_occlusion.txt'
        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_occu_file), 'r') as f:
            annos_occu_test = f.readlines()
        names_occu = [x.strip().split() for x in annos_occu_test]
        names_occu = [x[0]+'_'+x[-1] for x in names_occu]
        with open(os.path.join(root_folder, 'WFLW', 'test_occu.txt'), 'w') as f:
            for name_occu in names_occu:
                if name_occu in names_mapping:
                    image_crop_name, anno = names_mapping[name_occu]
                    f.write(image_crop_name+' ')
                    for x,y in anno:
                        f.write(str(x)+' '+str(y)+' ')
                    f.write('\n')
                else:
                    print('error!')
                    exit(0)


        test_blur_file = 'list_98pt_test_blur.txt'
        with open(os.path.join(root_folder, 'WFLW', 'WFLW_annotations', 'list_98pt_test', test_blur_file), 'r') as f:
            annos_blur_test = f.readlines()
        names_blur = [x.strip().split() for x in annos_blur_test]
        names_blur = [x[0]+'_'+x[-1] for x in names_blur]
        with open(os.path.join(root_folder, 'WFLW', 'test_blur.txt'), 'w') as f:
            for name_blur in names_blur:
                if name_blur in names_mapping:
                    image_crop_name, anno = names_mapping[name_blur]
                    f.write(image_crop_name+' ')
                    for x,y in anno:
                        f.write(str(x)+' '+str(y)+' ')
                    f.write('\n')
                else:
                    print('error!')
                    exit(0)
        gen_meanface(root_folder, data_name)
    ################################################################################################################
    elif data_name == 'AFLW':
        mat = hdf5storage.loadmat('../data/AFLW/AFLWinfo_release.mat')
        bboxes = mat['bbox']
        annos = mat['data']
        mask_new = mat['mask_new']
        nameList = mat['nameList']
        ra = mat['ra'][0]
        train_indices = ra[:20000]
        test_indices = ra[20000:]

        with open(os.path.join(root_folder, 'AFLW', 'train.txt'), 'w') as f:
            for index in train_indices:
                # from matlab index
                image_name = nameList[index-1][0][0]
                bbox = bboxes[index-1]
                anno = annos[index-1]
                image_crop, anno = process_aflw(root_folder, image_name, bbox, anno, target_size)
                pad_num = 5-len(str(index))
                image_crop_name = 'aflw_train_' + '0' * pad_num + str(index) + '.jpg'
                print(image_crop_name)
                cv2.imwrite(os.path.join(root_folder, 'AFLW', 'images_train', image_crop_name), image_crop)
                f.write(image_crop_name+' ')
                for x,y in anno:
                    f.write(str(x)+' '+str(y)+' ')
                f.write('\n')

        with open(os.path.join(root_folder, 'AFLW', 'test.txt'), 'w') as f:
            for index in test_indices:
                # from matlab index
                image_name = nameList[index-1][0][0]
                bbox = bboxes[index-1]
                anno = annos[index-1]
                image_crop, anno = process_aflw(root_folder, image_name, bbox, anno, target_size)
                pad_num = 5-len(str(index))
                image_crop_name = 'aflw_test_' + '0' * pad_num + str(index) + '.jpg'
                print(image_crop_name)
                cv2.imwrite(os.path.join(root_folder, 'AFLW', 'images_test', image_crop_name), image_crop)
                f.write(image_crop_name+' ')
                for x,y in anno:
                    f.write(str(x)+' '+str(y)+' ')
                f.write('\n')
        gen_meanface(root_folder, data_name)
    else:
        print('Wrong data!')

if __name__ == '__main__':
    if len(sys.argv) < 2:
        print('please input the data name.')
        print('1. data_300W')
        print('2. COFW')
        print('3. WFLW')
        print('4. AFLW')
        exit(0)
    else:
        data_name = sys.argv[1]
        gen_data('../data', data_name, 256)