| """Pre-process a dataset for model training: resize & center-crop every image |
| into a bucket within [min, max] on the long side, aligned to a step. |
| Images are written as sequentially numbered files with paired .txt captions. |
| Preserves original format (JPG q97, PNG lossless). |
| |
| Usage: |
| python preprocess_dataset.py --input /path/to/img --output /path/to/out |
| python preprocess_dataset.py --input /path/to/img --output /path/to/out --min-size 1024 --max-size 1152 --step 64 |
| python preprocess_dataset.py --input /path/to/img --output /path/to/out --dry-run |
| """ |
|
|
| import argparse |
| import os |
| import struct |
| from PIL import Image, ImageFile |
|
|
| ImageFile.LOAD_TRUNCATED_IMAGES = True |
| from tqdm import tqdm |
|
|
|
|
| JPEG_QUALITY = 97 |
| NAME_WIDTH = 7 |
|
|
| IMAGE_EXTS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'} |
|
|
|
|
| def process_image(img, min_size, max_size, step): |
| """Resize + crop to fit [min_size, max_size] with step alignment.""" |
| w, h = img.size |
| if max(w, h) > max_size: |
| if w >= h: |
| new_w, new_h = max_size, int(max_size * h / w) |
| else: |
| new_w, new_h = int(max_size * w / h), max_size |
| else: |
| new_w, new_h = w, h |
| if min(new_w, new_h) < min_size: |
| if new_w <= new_h: |
| new_w, new_h = min_size, int(min_size * new_h / new_w) |
| else: |
| new_w, new_h = int(min_size * new_w / new_h), min_size |
| crop_w = min(max_size, (new_w // step) * step) |
| crop_h = min(max_size, (new_h // step) * step) |
| crop_w = max(min_size, crop_w) |
| crop_h = max(min_size, crop_h) |
| img = img.convert("RGB").resize((new_w, new_h), Image.LANCZOS) |
| left = (new_w - crop_w) // 2 |
| top = (new_h - crop_h) // 2 |
| return img.crop((left, top, left + crop_w, top + crop_h)) |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Resize dataset into buckets") |
| parser.add_argument("--input", required=True) |
| parser.add_argument("--output", required=True) |
| parser.add_argument("--min-size", type=int, default=768, help="bucket min side") |
| parser.add_argument("--max-size", type=int, default=1280, help="bucket max side") |
| parser.add_argument("--step", type=int, default=16, help="bucket alignment step") |
| parser.add_argument("--dry-run", action="store_true") |
| args = parser.parse_args() |
|
|
| os.makedirs(args.output, exist_ok=True) |
|
|
| |
| print("Scanning recursively...") |
| captions = {} |
| image_paths = [] |
|
|
| for root, dirs, files in os.walk(args.input): |
| txt_stems = set() |
| for fname in files: |
| if fname.lower().endswith('.txt'): |
| stem = os.path.join(root, os.path.splitext(fname)[0]) |
| txt_stems.add(stem) |
| with open(os.path.join(root, fname)) as f: |
| captions[stem] = f.read() |
|
|
| for fname in sorted(files): |
| if os.path.splitext(fname)[1].lower() in IMAGE_EXTS: |
| stem = os.path.join(root, os.path.splitext(fname)[0]) |
| if stem in txt_stems: |
| image_paths.append(os.path.join(root, fname)) |
|
|
| image_paths.sort() |
| print(f"Found {len(image_paths)} images, {len(captions)} paired .txt captions") |
| print(f"Min: {args.min_size}, Max: {args.max_size}, Step: {args.step}, Quality: {JPEG_QUALITY}") |
|
|
| if args.dry_run: |
| for path in tqdm(image_paths, desc="Dry run"): |
| try: |
| img = Image.open(path) |
| img.load() |
| out = process_image(img, args.min_size, args.max_size, args.step) |
| print(f"{path}: {img.size} -> {out.size}") |
| except (OSError, IOError, struct.error) as e: |
| print(f"{path}: SKIP ({e})") |
| return |
|
|
| stats = {"upscaled": 0, "downscaled": 0, "unchanged": 0, "skipped": 0} |
|
|
| for idx, src_path in enumerate(tqdm(image_paths, desc="Processing"), start=1): |
| stem_no_ext = os.path.splitext(src_path)[0] |
| ext = os.path.splitext(src_path)[1].lower() |
|
|
| try: |
| img = Image.open(src_path) |
| img.load() |
| orig_w, orig_h = img.size |
| out = process_image(img, args.min_size, args.max_size, args.step) |
| except (OSError, IOError, struct.error) as e: |
| tqdm.write(f" Skip {os.path.basename(src_path)}: {e}") |
| stats["skipped"] += 1 |
| continue |
|
|
| if out.size == (orig_w, orig_h): |
| stats["unchanged"] += 1 |
| elif max(out.size) > max(orig_w, orig_h): |
| stats["upscaled"] += 1 |
| else: |
| stats["downscaled"] += 1 |
|
|
| |
| name = f"{idx:0{NAME_WIDTH}d}" |
| out_path = os.path.join(args.output, f"{name}{ext}") |
| if ext in {'.jpg', '.jpeg'}: |
| out.save(out_path, quality=JPEG_QUALITY) |
| else: |
| out.save(out_path) |
|
|
| |
| if stem_no_ext in captions: |
| txt_path = os.path.join(args.output, f"{name}.txt") |
| with open(txt_path, "w") as f: |
| f.write(captions[stem_no_ext]) |
|
|
| print(f"\nDone: {stats['unchanged']} unchanged, " |
| f"{stats['upscaled']} upscaled, " |
| f"{stats['downscaled']} downscaled, " |
| f"{stats['skipped']} skipped") |
| print(f"Output: {args.output}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
