Hugging Face's logo

recoilme
/
img

Model card Files
xet
 Community
img
File size: 4,732 Bytes
ffd809f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
"""
Pre-process dataset for Z-Image LoRA training.
Reads images from --input, resizes to [768, 1280] step 16,
writes as sequentially numbered JPEGs + paired .txt captions to --output.

Usage:
  python preprocess_dataset.py --input /path/to/img --output /path/to/img3
"""

import argparse
import os
import struct
from PIL import Image, ImageFile

# Allow truncated images
ImageFile.LOAD_TRUNCATED_IMAGES = True
from tqdm import tqdm


MIN_SIZE = 768
MAX_SIZE = 1280
STEP = 16
JPEG_QUALITY = 97
NAME_WIDTH = 7  # 0000001, 0000002, ...

IMAGE_EXTS = {'.jpg', '.jpeg', '.png', '.webp', '.bmp'}


def process_image(img: Image.Image) -> Image.Image:
    """Resize + crop to fit [MIN_SIZE, MAX_SIZE] with STEP alignment."""
    w, h = img.size

    # 1. Fit max side to MAX_SIZE (only downsample if too big)
    if max(w, h) > MAX_SIZE:
        if w >= h:
            new_w = MAX_SIZE
            new_h = int(MAX_SIZE * h / w)
        else:
            new_h = MAX_SIZE
            new_w = int(MAX_SIZE * w / h)
    else:
        new_w, new_h = w, h

    # 2. If min side is still too small, fit to MIN_SIZE
    if min(new_w, new_h) < MIN_SIZE:
        if new_w <= new_h:
            new_w = MIN_SIZE
            new_h = int(MIN_SIZE * new_h / new_w)
        else:
            new_h = MIN_SIZE
            new_w = int(MIN_SIZE * new_w / new_h)

    # 3. Crop to STEP-aligned, clamped to [MIN_SIZE, MAX_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)

    # 4. Resize + center crop
    img = img.convert("RGB").resize((new_w, new_h), Image.LANCZOS)
    left = (new_w - crop_w) // 2
    top = (new_h - crop_h) // 2
    img = img.crop((left, top, left + crop_w, top + crop_h))

    return img


def main():
    parser = argparse.ArgumentParser(description="Resize dataset for Z-Image")
    parser.add_argument("--input", required=True)
    parser.add_argument("--output", required=True)
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args()

    os.makedirs(args.output, exist_ok=True)

    # Collect image files and find paired .txt captions
    all_files = sorted(os.listdir(args.input))
    image_files = [f for f in all_files if os.path.splitext(f)[1].lower() in IMAGE_EXTS]

    # Build mapping: basename -> caption text (if .txt exists with same basename)
    captions = {}
    for f in all_files:
        if f.lower().endswith(".txt"):
            stem = os.path.splitext(f)[0]
            with open(os.path.join(args.input, f)) as fh:
                captions[stem] = fh.read()

    print(f"Found {len(image_files)} images, {len(captions)} paired .txt captions")
    print(f"Min: {MIN_SIZE}, Max: {MAX_SIZE}, Step: {STEP}, Quality: {JPEG_QUALITY}")

    if args.dry_run:
        for fname in tqdm(image_files, desc="Dry run"):
            try:
                img = Image.open(os.path.join(args.input, fname))
                img.load()
                out = process_image(img)
                print(f"{fname}: {img.size} -> {out.size}")
            except (OSError, IOError, struct.error) as e:
                print(f"{fname}: SKIP ({e})")
        return

    stats = {"upscaled": 0, "downscaled": 0, "unchanged": 0, "skipped": 0}

    for idx, fname in enumerate(tqdm(image_files, desc="Processing"), start=1):
        src_path = os.path.join(args.input, fname)
        stem = os.path.splitext(fname)[0]

        try:
            img = Image.open(src_path)
            img.load()  # force load to catch truncated images early
            orig_w, orig_h = img.size
            out = process_image(img)
        except (OSError, IOError, struct.error) as e:
            tqdm.write(f"  Skip {fname}: {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

        # Sequential output names: 0000001.jpg, 0000001.txt
        name = f"{idx:0{NAME_WIDTH}d}"
        out_path = os.path.join(args.output, name + ".jpg")
        out.save(out_path, quality=JPEG_QUALITY)

        # Copy caption if paired .txt exists
        if stem in captions:
            txt_path = os.path.join(args.output, name + ".txt")
            with open(txt_path, "w") as f:
                f.write(captions[stem])

    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()