import os
from beer.data.tools import check_xml_and_img_file
from beer.data.tools import parse_args
from utils.cropper import ImageListCropper
from utils.io import create_map_pdtxt
from utils.io import read_text_file
from utils.io import read_label_as_list
from utils.io import get_label_list_from_category_index
from utils.list import create_file_list
from utils.list import create_train_val_list
def process_all(lists, output_root, split='&!&'):
    if not os.path.exists(output_root):
        os.makedirs(output_root)
    for count, paths in enumerate(lists):
        print(paths)
        img_path, xml_path = paths.split(split)
        out_root = os.path.join(output_root, '{:04}'.format(count // 1000),
                                '{:08}'.format(count))
        if not os.path.exists(out_root):
            os.makedirs(out_root)
        print(img_path)
        cropper = ImageListCropper(img_path.strip(), xml_path.strip(), out_root,
                                   cropped_size=(800,800 ), stride=(600, 600))
        cropper.update(class_names, output_root + '/break.txt')
def _make_data(prefix, postfix, split='&!&'):
    train = 'train{}'.format(postfix)
    val = 'val{}'.format(postfix)
    origin_data = os.path.join(args.root_path, args.dataset)
    output_data = os.path.join(args.root_path, args.target)
    file_list, _ = create_file_list(origin_data,
                                    filtering=check_xml_and_img_file,
                                    params=[class_names, split])
    print(file_list)
    create_train_val_list(file_list, args.root_path, prefix, args.postfix)
    train_list = read_text_file(os.path.join(args.root_path, '{}{}.txt'.format(prefix, train)))
    train_path = os.path.join(output_data, train)
    process_all(train_list, train_path, split)
    create_file_list(train_path, os.path.join(args.root_path, '{}.txt'.format(train)),
                     filtering=check_xml_and_img_file,
                     params=[class_names, split])
    val_list = read_text_file(os.path.join(args.root_path, '{}{}.txt'.format(prefix, val)))
    val_path = os.path.join(output_data, val)
    process_all(val_list, val_path, split)
    create_file_list(val_path, os.path.join(args.root_path, '{}.txt'.format(val)),
                     filtering=check_xml_and_img_file,
                     params=[class_names, split])
    create_map_pdtxt(os.path.join(args.root_path, '{}_{}_map.pdtxt'.format(prefix, postfix)), class_names)


if __name__ == '__main__':
    args = parse_args()
    # class_names = read_label_as_list(args.label_file, args.class_num, args.instance)
    # class_names = get_label_list_from_category_index(pd_file=args.label_file, class_num=args.class_num)
    with open(args.label_file, 'rb') as file:
        class_names = file.readlines()
        class_names = list(map(lambda x: x.strip().decode(), class_names))
        # print(class_names)
    _make_data(args.prefix, args.postfix)