Switch to args instead of hardcoding variables

preprocess_images.py

@@ -1,23 +1,22 @@
 import os
 import hashlib
+import argparse
 from tqdm import tqdm
 from PIL import Image
 from queue import Queue
 from threading import Thread
 
-# target resolution [latent res * 8]
-resolution = 768
-# threads used for resizing
-threads = 4
-# source folder with images
-folder = "raw"
-
 if not os.path.isdir("images"):
 	os.mkdir("images")
 
-def process(fname):
-	global folder
-	global resolution
+def parse_args():
+	parser = argparse.ArgumentParser(description="Preprocess images")
+	parser.add_argument("-r", "--res", type=int, default=768, help="Target resolution")
+	parser.add_argument("-t", "--threads", type=int, default=4, help="No. of CPU threads to use")
+	parser.add_argument('--src', default="raw", help="Source folder with images")
+	return parser.parse_args()
+
+def process(fname, folder, resolution):
 	src = os.path.join(folder, fname)
 	md5 = hashlib.md5(open(src,'rb').read()).hexdigest()
 	out = os.path.join("images", f"{md5}.png")
@@ -28,22 +27,28 @@ def process(fname):
 	img = img.resize((resolution,resolution), Image.LANCZOS)
 	img.save(out)
 
-def thread(queue, pbar):
+def thread(queue, pbar, folder, resolution):
 	while not queue.empty():
 		fname = queue.get()
-		process(fname)
+		process(fname, folder, resolution)
 		queue.task_done()
 		pbar.update()
 
-files = os.listdir(folder)
+args = parse_args()
+files = os.listdir(args.src)
 pbar = tqdm(total=len(files),unit="img")
 queue = Queue()
 [queue.put(x) for x in files]
 
-for _ in range(threads):
+for _ in range(args.threads):
 	Thread(
 		target=thread,
-		args=(queue,pbar),
+		args=(
+			queue,
+			pbar,
+			args.src,
+			args.res,
+		),
 		daemon=True,
 	).start()