<reponame>duanzhiihao/mycv
import os
from argparse import ArgumentParser
from tqdm import tqdm
from PIL import Image
# from multiprocessing import Pool


alg_dict = {
    'lanczos': Image.LANCZOS,
    'nearest': Image.NEAREST,
    'bilinear': Image.BILINEAR,
    'bicubic': Image.BICUBIC,
    'hamming': Image.HAMMING,
    'box': Image.BOX
}


def parse_arguments():
    parser = ArgumentParser()
    parser.add_argument('-i', '--in_dir', help="Input directory with source images", required=True)
    parser.add_argument('-o', '--out_dir', help="Output directory for resized images", required=True)
    parser.add_argument('-s', '--size', help="Size of an output image (e.g. 32 results in (32x32) image)",
                        default=32, type=int)
    parser.add_argument('-a', '--algorithm', help="Algorithm used for resampling: lanczos, nearest,"
                                                  " bilinear, bicubic, box, hamming",
                        default='lanczos')

    parser.add_argument('-f', '--full', help="Use all algorithms, create subdirectory for each algorithm output",
                        action='store_true')
    parser.add_argument('-j', '--workers', help="Number of sub-processes that run different folders "
                                                  "in the same time ",
                        default=1, type=int)
    args = parser.parse_args()

    return args.in_dir, args.out_dir, args.algorithm, args.size, \
           args.full, args.workers


def str2alg(str):
    str = str.lower()
    return alg_dict.get(str, None)


# Takes in_dir, out_dir and alg as strings
# resize images from in_dir using algorithm deduced from
# alg string and puts them to "out_dir/alg/" folder
def resize_img_folder(in_dir, out_dir, tgt_size, alg):
    alg = str2alg(alg)

    if alg is None:
        print("Sorry but this algorithm (%s) is not available, use help for more info." % alg)
        return

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

    for filename in os.listdir(in_dir):
        im = Image.open(os.path.join(in_dir, filename))

        # Convert grayscale images into 3 channels
        if im.mode != "RGB":
            im = im.convert(mode="RGB")

        im_resized = im.resize((tgt_size, tgt_size), alg)
        # Get rid of extension (.jpg or other)
        filename = os.path.splitext(filename)[0]
        im_resized.save(os.path.join(out_dir, filename + '.png'))


def main():
    in_dir, out_root, alg_str, size, full, workers = parse_arguments()
    assert workers == 1

    print('Starting ...')

    if full is False:
        algorithms = [alg_str]
    else:
        algorithms = alg_dict.keys()

    for alg in algorithms:
        print('Using algorithm %s ...' % alg)
        cls_names = [name for name in os.listdir(in_dir)]
        pbar = tqdm(cls_names)
        for i, cname in enumerate(pbar):
            src = os.path.join(in_dir, cname)
            tgt = os.path.join(out_root, cname)
            pbar.set_description(f'Saving to {tgt}')
            resize_img_folder(src, tgt, size, alg)
    print("Finished.")


if __name__ == '__main__':
    main()

import os
from argparse import ArgumentParser
from tqdm import tqdm
from PIL import Image
# from multiprocessing import Pool


alg_dict = {
    'lanczos': Image.LANCZOS,
    'nearest': Image.NEAREST,
    'bilinear': Image.BILINEAR,
    'bicubic': Image.BICUBIC,
    'hamming': Image.HAMMING,
    'box': Image.BOX
}


def parse_arguments():
    parser = ArgumentParser()
    parser.add_argument('-i', '--in_dir', help="Input directory with source images", required=True)
    parser.add_argument('-o', '--out_dir', help="Output directory for resized images", required=True)
    parser.add_argument('-s', '--size', help="Size of an output image (e.g. 32 results in (32x32) image)",
                        default=32, type=int)
    parser.add_argument('-a', '--algorithm', help="Algorithm used for resampling: lanczos, nearest,"
                                                  " bilinear, bicubic, box, hamming",
                        default='lanczos')

    parser.add_argument('-f', '--full', help="Use all algorithms, create subdirectory for each algorithm output",
                        action='store_true')
    parser.add_argument('-j', '--workers', help="Number of sub-processes that run different folders "
                                                  "in the same time ",
                        default=1, type=int)
    args = parser.parse_args()

    return args.in_dir, args.out_dir, args.algorithm, args.size, \
           args.full, args.workers


def str2alg(str):
    str = str.lower()
    return alg_dict.get(str, None)


# Takes in_dir, out_dir and alg as strings
# resize images from in_dir using algorithm deduced from
# alg string and puts them to "out_dir/alg/" folder
def resize_img_folder(in_dir, out_dir, tgt_size, alg):
    alg = str2alg(alg)

    if alg is None:
        print("Sorry but this algorithm (%s) is not available, use help for more info." % alg)
        return

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

    for filename in os.listdir(in_dir):
        im = Image.open(os.path.join(in_dir, filename))

        # Convert grayscale images into 3 channels
        if im.mode != "RGB":
            im = im.convert(mode="RGB")

        im_resized = im.resize((tgt_size, tgt_size), alg)
        # Get rid of extension (.jpg or other)
        filename = os.path.splitext(filename)[0]
        im_resized.save(os.path.join(out_dir, filename + '.png'))


def main():
    in_dir, out_root, alg_str, size, full, workers = parse_arguments()
    assert workers == 1

    print('Starting ...')

    if full is False:
        algorithms = [alg_str]
    else:
        algorithms = alg_dict.keys()

    for alg in algorithms:
        print('Using algorithm %s ...' % alg)
        cls_names = [name for name in os.listdir(in_dir)]
        pbar = tqdm(cls_names)
        for i, cname in enumerate(pbar):
            src = os.path.join(in_dir, cname)
            tgt = os.path.join(out_root, cname)
            pbar.set_description(f'Saving to {tgt}')
            resize_img_folder(src, tgt, size, alg)
    print("Finished.")


if __name__ == '__main__':
    main()
# from multiprocessing import Pool # Takes in_dir, out_dir and alg as strings # resize images from in_dir using algorithm deduced from # alg string and puts them to "out_dir/alg/" folder # Convert grayscale images into 3 channels # Get rid of extension (.jpg or other)