convert_seg.py 文件源码

def main(directory, convert_directory, test, crop_size, extension):
    try:
        os.mkdir(convert_directory)
    except OSError:
        pass

    supported_extensions = set(['jpg', 'png', 'tiff', 'jpeg', 'tif'])

    # filenames = [os.path.join(dp, f) for dp, dn, fn in os.walk(directory)
    # for f in fn if f.split('.')[-1].lower() in supported_extensions]
    filenames = [each for each in os.listdir(
        directory) if each.endswith('.jpg')]
    filenames = [os.path.join(directory, filename.strip(
        '\n')) for filename in filenames]
    # with open('/home/artelus_server/data/segment_artelus/train.txt', 'r') as f:
    #    filenames = f.readlines()
    # filenames = [os.path.join(directory, filename.strip(
    #    '\n') + '.jpg') for filename in filenames]
    filenames = sorted(filenames)

    if test:
        names = data.get_names(filenames)
        y = data.get_labels(names)
        for f, level in zip(filenames, y):
            if level == 1:
                try:
                    img = convert(f, crop_size)
                    img.show()
                    Image.open(f).show()
                    real_raw_input = vars(__builtins__).get('raw_input', input)
                    real_raw_input('enter for next')
                except KeyboardInterrupt:
                    exit(0)

    print("Resizing images in {} to {}, this takes a while."
          "".format(directory, convert_directory))

    n = len(filenames)
    # process in batches, sometimes weird things happen with Pool on my machine
    batchsize = 500
    batches = n // batchsize + 1
    pool = Pool(N_PROC)

    args = []

    for f in filenames:
        label_f = f[:-4] + '_final_mask.png'
        args.append((convert, (directory, convert_directory, f, label_f, crop_size,
                               extension)))

    for i in range(batches):
        print("batch {:>2} / {}".format(i + 1, batches))
        pool.map(process, args[i * batchsize: (i + 1) * batchsize])

    pool.close()

    print('done')