resize_img_dir.py

'''
Resizes images in source image directory within given size bounds (keeping 
aspect ratio) and outputs in target directory with identical directory tree
structure. Uses Magick for image resizing.
'''
import os
import argparse
import subprocess
from pathlib import Path

from PIL import Image
from tqdm import tqdm
from pqdm.processes import pqdm


def collect_sources(curr_dir, sources, extension):
	for item in curr_dir.iterdir():
		if item.is_dir():
			collect_sources(item)
		elif item.is_file() and item.suffix == extension:
			sources.append(item)


def resolve_targets(sources, source_dir, target_dir):
	source_dir_str = str(source_dir)
	target_dir_str = str(target_dir)
	return [Path(str(s).replace(source_dir_str,target_dir_str)) for s in sources]


def get_resize_command(source, target, size, force):
	op = '' if force else '\\>' # `\>` operator ensures that images smaller than size will merely be copied over instead of upsampled 
	return ['magick', str(source.absolute()), '-resize', f'{size}x{size}{op}', str(target.absolute())]


def resize(command):
	return subprocess.run(command, text=True, capture_output=True)


def verify_targets(targets, size):
	for tgt in tqdm(targets):
		if tgt.is_file():
			img = Image.open(tgt)
			img.verify()
			width, height = img.size
			img.close()
			if width > size or height > size:
				print(f'IMAGE SIZE {width}x{height} EXCEEDED!', tgt)
		else:
			print('FILE NOT FOUND!', tgt)


def parse_args():
	parser = argparse.ArgumentParser(description='.')
	parser.add_argument('--size',       '-s',   type=int,	default=512,                  help='Target size x size.')
	parser.add_argument('--source_dir',	'-sd',  type=str, default='images',             help='Source images directory.')
	parser.add_argument('--target_dir', '-td',  type=str,	default='images_resized_512',	help='Target images directory.')
	parser.add_argument('--extension',	'-ext',	type=str,	default='.jpg',								help='The image file extension.')
	parser.add_argument('--force',			'-f',		action='store_true',										help='Whether to force image to size. If not indicated, images smaller than size bounds will be copied over instead of upsampled.')
	args = parser.parse_args()
	args.source_dir = Path(args.source_dir)
	args.target_dir = Path(args.target_dir)
	args.target_dir.mkdir(exist_ok=True)
	return args


def main():
	args = parse_args()
	sources = []
	print('Collecting sources...')
	collect_sources(args.source_dir, sources, args.extension)
	print(f'{len(sources)} sources collected!')
	targets = resolve_targets(sources, args.source_dir, args.target_dir)

	print(f'Resizing images within size {args.size}x{args.size}...')  
	args_list = [get_resize_command(sources[i], targets[i], args.size, args.force) for i in range(len(sources))]
	results = pqdm(args_list, resize, n_jobs=os.cpu_count())
	for r in results:
		print(r.stdout)
		print(r.stderr)

	print('Verifying targets...')
	verify_targets(targets, args.size)

if __name__ == '__main__': main()