Add a control generation script.
This commit is contained in:
Jaret Burkett
256
scripts/add_mask_dataset.py
Normal file
256
scripts/add_mask_dataset.py
Normal file
@@ -0,0 +1,256 @@
|
||||
import os
|
||||
import sys
|
||||
import queue
|
||||
import random
|
||||
import argparse
|
||||
import threading
|
||||
import traceback
|
||||
|
||||
import torch
|
||||
from tqdm import tqdm
|
||||
|
||||
# allow importing from project root
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from toolkit.control_generator import ControlGenerator, img_ext_list
|
||||
|
||||
|
||||
def control_exists(img_path, control_type):
|
||||
# mirrors the lookup in ControlGenerator.get_control_path so we can skip
|
||||
# images another instance has already finished
|
||||
controls_folder = os.path.join(os.path.dirname(img_path), "_controls")
|
||||
file_name_no_ext = os.path.splitext(os.path.basename(img_path))[0]
|
||||
file_name_no_ext_control = f"{file_name_no_ext}.{control_type}"
|
||||
for ext in img_ext_list:
|
||||
if os.path.exists(os.path.join(controls_folder, file_name_no_ext_control + ext)):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
# sentinel pushed through the queues to tell workers to stop
|
||||
_DONE = object()
|
||||
|
||||
|
||||
def run_pipeline(control_gen, img_list, control_type, regen, n_load, n_save):
|
||||
# Three-stage pipeline so the GPU never waits on disk/CPU work:
|
||||
# loaders (N threads) -> read + exif + resize + preprocess to a CPU tensor
|
||||
# gpu worker (1 thread) -> model forward only (kept single so VRAM is bounded)
|
||||
# savers (M threads) -> postprocess (resize/alpha) + encode + write
|
||||
# The heavy CPU work (resize/normalize on input, resize/convert on output) is
|
||||
# pushed onto the loader/saver threads so the GPU thread does almost nothing
|
||||
# but the forward pass. Bounded queues apply backpressure so we don't load the
|
||||
# whole dataset into RAM.
|
||||
infer_q = queue.Queue(maxsize=n_load * 2)
|
||||
save_q = queue.Queue(maxsize=n_save * 2)
|
||||
path_q = queue.Queue()
|
||||
for img_path in img_list:
|
||||
path_q.put(img_path)
|
||||
|
||||
# miniters=1 disables tqdm's dynamic-miniters heuristic (which otherwise
|
||||
# raises the redraw threshold after a fast burst and makes the bar look
|
||||
# frozen); mininterval keeps redraws time-based and cheap.
|
||||
pbar = tqdm(total=len(img_list), desc=f"Generating {control_type}",
|
||||
miniters=1, mininterval=0.5)
|
||||
pbar_lock = threading.Lock()
|
||||
# set on completion OR on Ctrl-C; every blocking call below uses a timeout and
|
||||
# re-checks this so the worker threads can actually be shut down.
|
||||
stop_event = threading.Event()
|
||||
|
||||
def put(q, item):
|
||||
# interruptible put: blocks until there's room, but wakes periodically so
|
||||
# a stop request (or KeyboardInterrupt on the main thread) is honored.
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
q.put(item, timeout=0.2)
|
||||
return
|
||||
except queue.Full:
|
||||
continue
|
||||
|
||||
def loader():
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
img_path = path_q.get_nowait()
|
||||
except queue.Empty:
|
||||
break
|
||||
try:
|
||||
if not regen and control_exists(img_path, control_type):
|
||||
# another instance (or a previous run) already did it
|
||||
with pbar_lock:
|
||||
pbar.update(1)
|
||||
continue
|
||||
image = control_gen.load_image(img_path)
|
||||
payload = control_gen.preprocess(image, control_type)
|
||||
put(infer_q, (img_path, image, payload))
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
with pbar_lock:
|
||||
pbar.update(1)
|
||||
|
||||
def saver():
|
||||
while not stop_event.is_set():
|
||||
try:
|
||||
item = save_q.get(timeout=0.2)
|
||||
except queue.Empty:
|
||||
continue
|
||||
if item is _DONE:
|
||||
break
|
||||
img_path, image, result = item
|
||||
try:
|
||||
out_image = control_gen.postprocess(result, image, control_type)
|
||||
save_path = control_gen.control_save_path(img_path, control_type)
|
||||
control_gen.save_control(out_image, save_path)
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
finally:
|
||||
with pbar_lock:
|
||||
pbar.update(1)
|
||||
|
||||
loaders = [threading.Thread(target=loader, daemon=True) for _ in range(n_load)]
|
||||
savers = [threading.Thread(target=saver, daemon=True) for _ in range(n_save)]
|
||||
for t in loaders + savers:
|
||||
t.start()
|
||||
|
||||
# GPU stage runs on the main thread: pull preprocessed tensors, run the
|
||||
# forward pass, hand the raw result off to the savers. We stop once every
|
||||
# loader has exited and nothing is left queued for inference.
|
||||
interrupted = False
|
||||
try:
|
||||
while not stop_event.is_set():
|
||||
if not any(t.is_alive() for t in loaders) and infer_q.empty():
|
||||
break
|
||||
try:
|
||||
img_path, image, payload = infer_q.get(timeout=0.1)
|
||||
except queue.Empty:
|
||||
continue
|
||||
try:
|
||||
result = control_gen.run_inference(payload, control_type)
|
||||
put(save_q, (img_path, image, result))
|
||||
except Exception:
|
||||
traceback.print_exc()
|
||||
with pbar_lock:
|
||||
pbar.update(1)
|
||||
except KeyboardInterrupt:
|
||||
interrupted = True
|
||||
print("\nInterrupted, shutting down...")
|
||||
|
||||
if interrupted:
|
||||
# abort: tell every worker to stop; pending items are dropped
|
||||
stop_event.set()
|
||||
else:
|
||||
# normal finish: let savers drain whatever is still queued, then stop
|
||||
for _ in savers:
|
||||
save_q.put(_DONE)
|
||||
|
||||
# join with a timeout so a stuck worker can never wedge shutdown; threads are
|
||||
# daemons, so anything still alive is torn down when we return.
|
||||
for t in savers:
|
||||
t.join(timeout=5)
|
||||
pbar.close()
|
||||
if interrupted:
|
||||
raise KeyboardInterrupt
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate masks for a dataset using the ControlGenerator"
|
||||
)
|
||||
parser.add_argument("img_dir", type=str, help="Path to image directory")
|
||||
parser.add_argument(
|
||||
"--control",
|
||||
type=str,
|
||||
default="mask",
|
||||
choices=["mask", "inpaint", "depth", "pose", "line", "sapiens2_mask"],
|
||||
help="Control type to generate (default: mask)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--device", type=str, default="cuda", help="Device to run on (default: cuda)"
|
||||
)
|
||||
parser.add_argument("--debug", action="store_true", help="Enable debug mode")
|
||||
parser.add_argument(
|
||||
"--regen",
|
||||
action="store_true",
|
||||
help="Regenerate controls even if they already exist",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shuffle",
|
||||
action="store_true",
|
||||
help="Shuffle image order so multiple instances on the same dataset "
|
||||
"don't chase the same images",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--load-workers",
|
||||
type=int,
|
||||
default=16,
|
||||
help="Number of threads for loading/resizing images (default: 4)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-workers",
|
||||
type=int,
|
||||
default=16,
|
||||
help="Number of threads for saving controls (default: 4)",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
img_dir = args.img_dir
|
||||
if not os.path.isdir(img_dir):
|
||||
print(f"Error: {img_dir} is not a directory")
|
||||
sys.exit(1)
|
||||
|
||||
# find images, skipping existing _controls folders and hidden files
|
||||
img_list = []
|
||||
for root, dirs, files in os.walk(img_dir):
|
||||
if "_controls" in root:
|
||||
continue
|
||||
for file in files:
|
||||
if file.startswith("."):
|
||||
continue
|
||||
if file.lower().endswith(tuple(img_ext_list)):
|
||||
img_list.append(os.path.join(root, file))
|
||||
|
||||
if len(img_list) == 0:
|
||||
print(f"Error: no images found in {img_dir}")
|
||||
sys.exit(1)
|
||||
|
||||
# filter out images that already have controls up front so the progress bar
|
||||
# reflects only real work (otherwise it races through thousands of skips and
|
||||
# the rate/ETA are meaningless). The loader still re-checks just before
|
||||
# processing to handle the multi-instance race.
|
||||
if not args.regen:
|
||||
total = len(img_list)
|
||||
img_list = [p for p in img_list if not control_exists(p, args.control)]
|
||||
skipped = total - len(img_list)
|
||||
if skipped:
|
||||
print(f"Skipping {skipped} images that already have '{args.control}' controls")
|
||||
if len(img_list) == 0:
|
||||
print("All images already have controls. Nothing to do.")
|
||||
return
|
||||
|
||||
if args.shuffle:
|
||||
random.shuffle(img_list)
|
||||
|
||||
control_gen = ControlGenerator(torch.device(args.device))
|
||||
control_gen.debug = args.debug
|
||||
control_gen.regen = args.regen
|
||||
|
||||
interrupted = False
|
||||
try:
|
||||
run_pipeline(
|
||||
control_gen,
|
||||
img_list,
|
||||
args.control,
|
||||
args.regen,
|
||||
max(1, args.load_workers),
|
||||
max(1, args.save_workers),
|
||||
)
|
||||
except KeyboardInterrupt:
|
||||
interrupted = True
|
||||
finally:
|
||||
control_gen.cleanup()
|
||||
|
||||
if interrupted:
|
||||
sys.exit(130)
|
||||
print("Done")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user