import argparse import glob import json import logging import os import random import subprocess from io import BytesIO import pyarrow as pa import pyarrow.parquet as pq from PIL import Image from tools.dataset import BackgroundDataset, BackgroundIterableDataset logger = logging.getLogger(__name__) def iter_samples(dataset, streaming): if streaming: for sample in dataset: yield sample else: for idx in range(len(dataset)): yield dataset[idx] def parse_args(): parser = argparse.ArgumentParser(description="Sample background images for SynLayers.") parser.add_argument("--dataset-name", default="laion/laion2B-en-aesthetic") parser.add_argument( "--data-files", default="/project/llmsvgen/share/data/kmw_layered_dataset/laion2B-en-aesthetic-image/*.parquet", help="Parquet glob or list file.", ) parser.add_argument("--split", default="train") parser.add_argument("--cache-dir", default=None) parser.add_argument("--url-column", default="URL") parser.add_argument("--text-column", default="TEXT") parser.add_argument("--hash-column", default="hash") parser.add_argument( "--image-root", default="/project/llmsvgen/share/data/kmw_layered_dataset/laion2B-en-aesthetic-image", help="Local directory with downloaded images named by hash.", ) parser.add_argument( "--image-extensions", default=".jpg,.png,.jpeg,.webp", help="Comma-separated extensions to try for local images.", ) parser.add_argument("--image-size", type=int, default=None) parser.add_argument("--count", type=int, default=10) parser.add_argument("--streaming", action="store_true") parser.add_argument("--output-dir", default="./outputs/backgrounds") parser.add_argument( "--save-images", action="store_true", help="Save images if found in image-root.", ) parser.add_argument( "--download", action="store_true", help="Download a subset into image-root using img2dataset.", ) parser.add_argument( "--download-mode", choices=["auto", "img2dataset", "embedded"], default="auto", help="Download mode: auto-detect URL vs embedded bytes.", ) parser.add_argument("--processes", type=int, default=8) parser.add_argument("--threads", type=int, default=32) parser.add_argument("--resize", type=int, default=512) parser.add_argument("--build-splits", action="store_true") parser.add_argument("--train-count", type=int, default=19000) parser.add_argument("--val-count", type=int, default=1000) parser.add_argument("--test-count", type=int, default=200) parser.add_argument( "--skip-existing", action="store_true", help="Skip downloading/extracting images that already exist in image-root.", ) parser.add_argument( "--progress-interval", type=int, default=500, help="Log progress every N extracted images.", ) parser.add_argument( "--embedded-image-column", default="whole_image", help="Struct column containing embedded image bytes.", ) parser.add_argument( "--embedded-image-columns", default=None, help="Comma-separated embedded image columns to try in order.", ) parser.add_argument( "--embedded-image-bytes-key", default="bytes", help="Key inside embedded image struct that stores raw bytes.", ) parser.add_argument( "--embedded-image-path-key", default="path", help="Key inside embedded image struct that stores a path (if any).", ) parser.add_argument( "--embedded-caption-column", default="whole_caption", help="Caption column for embedded images.", ) parser.add_argument( "--embedded-id-column", default="id", help="ID column for embedded images.", ) parser.add_argument( "--size-multiple", type=int, default=8, help="Round width/height up to a multiple of this value.", ) parser.add_argument("--seed", type=int, default=42) parser.add_argument( "--sequential", action="store_true", help="Use dataset order instead of random sampling when building splits.", ) parser.add_argument( "--allow-partial", action="store_true", help="Allow writing splits even if there are not enough images.", ) parser.add_argument( "--id-as-path", action="store_true", help="Store image path in the id field instead of the raw key.", ) return parser.parse_args() def main(): logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s") args = parse_args() image_extensions = [ext.strip() for ext in args.image_extensions.split(",") if ext.strip()] if args.download: parquet_files = _expand_parquet_files(args.data_files) if not parquet_files: raise ValueError("No parquet files found. Check --data-files.") os.makedirs(args.image_root, exist_ok=True) download_mode = args.download_mode if args.embedded_image_columns: embedded_image_columns = [ col.strip() for col in args.embedded_image_columns.split(",") if col.strip() ] else: embedded_image_columns = [args.embedded_image_column] if download_mode == "auto": if _parquet_has_column(parquet_files, args.url_column): download_mode = "img2dataset" elif any( _parquet_has_column(parquet_files, col) for col in embedded_image_columns ): download_mode = "embedded" else: raise ValueError( "Could not detect download mode: missing URL and embedded image columns." ) if download_mode == "img2dataset": url_list_path = _prepare_download_parquet( parquet_files=parquet_files, output_dir=args.output_dir, count=args.count, seed=args.seed, url_column=args.url_column, text_column=args.text_column, hash_column=args.hash_column, ) cmd = [ "img2dataset", "--url_list", url_list_path, "--input_format", "parquet", "--url_col", args.url_column, "--caption_col", args.text_column, "--output_format", "files", "--output_folder", args.image_root, "--processes_count", str(args.processes), "--thread_count", str(args.threads), "--image_size", str(args.resize), "--resize_mode", "keep_ratio", ] logger.info("Downloading %d images into %s", args.count, args.image_root) subprocess.run(cmd, check=True) else: logger.info( "Extracting %d embedded images into %s", args.count, args.image_root, ) download_embedded_images( parquet_files=parquet_files, image_root=args.image_root, output_dir=args.output_dir, count=args.count, seed=args.seed, sequential=args.sequential, id_column=args.embedded_id_column, caption_column=args.embedded_caption_column, image_columns=embedded_image_columns, image_bytes_key=args.embedded_image_bytes_key, image_path_key=args.embedded_image_path_key, image_extensions=image_extensions, skip_existing=args.skip_existing, progress_interval=args.progress_interval, ) if args.build_splits: if _has_img2dataset_parquet(args.image_root): build_splits_from_img2dataset( image_root=args.image_root, output_dir=args.output_dir, train_count=args.train_count, val_count=args.val_count, test_count=args.test_count, seed=args.seed, sequential=args.sequential, allow_partial=args.allow_partial, id_as_path=args.id_as_path, image_extensions=image_extensions, size_multiple=args.size_multiple, ) else: build_splits( data_files=args.data_files, image_root=args.image_root, image_extensions=image_extensions, output_dir=args.output_dir, train_count=args.train_count, val_count=args.val_count, test_count=args.test_count, seed=args.seed, url_column=args.url_column, text_column=args.text_column, hash_column=args.hash_column, sequential=args.sequential, allow_partial=args.allow_partial, size_multiple=args.size_multiple, ) return if args.streaming: dataset = BackgroundIterableDataset( dataset_name=args.dataset_name, data_files=args.data_files, split=args.split, cache_dir=args.cache_dir, url_column=args.url_column, text_column=args.text_column, hash_column=args.hash_column, image_root=args.image_root, image_extensions=image_extensions, image_size=args.image_size, require_image=args.save_images, ) else: dataset = BackgroundDataset( dataset_name=args.dataset_name, data_files=args.data_files, split=args.split, cache_dir=args.cache_dir, url_column=args.url_column, text_column=args.text_column, hash_column=args.hash_column, image_root=args.image_root, image_extensions=image_extensions, image_size=args.image_size, max_items=args.count * 5, require_image=args.save_images, ) os.makedirs(args.output_dir, exist_ok=True) captions_path = os.path.join(args.output_dir, "captions.jsonl") saved = 0 with open(captions_path, "w", encoding="utf-8") as captions_file: for sample in iter_samples(dataset, args.streaming): image = sample.get("image") filename = None if args.save_images: if image is None: logger.warning("Skipping sample: local image not found.") continue filename = f"background_{saved:03d}.png" image.save(os.path.join(args.output_dir, filename)) captions_file.write( json.dumps( { "file": filename, "url": sample.get("url"), "text": sample.get("text"), "width": sample.get("width"), "height": sample.get("height"), "hash": sample.get("hash"), "aesthetic": sample.get("aesthetic"), "punsafe": sample.get("punsafe"), "pwatermark": sample.get("pwatermark"), }, ensure_ascii=False, ) + "\n" ) saved += 1 if saved >= args.count: break logger.info("Saved %d backgrounds to %s", saved, args.output_dir) def _expand_parquet_files(data_files): if isinstance(data_files, (list, tuple)): return list(data_files) if not data_files: return [] if os.path.exists(data_files) and data_files.endswith(".parquet"): return [data_files] return sorted(glob.glob(data_files)) def _parquet_has_column(parquet_files, column_name): if not column_name: return False for parquet_path in parquet_files: parquet_file = pq.ParquetFile(parquet_path) if column_name in parquet_file.schema.names: return True schema_arrow = getattr(parquet_file, "schema_arrow", None) if schema_arrow is not None and column_name in schema_arrow.names: return True return False def _has_img2dataset_parquet(image_root): if not image_root or not os.path.exists(image_root): return False return bool(glob.glob(os.path.join(image_root, "*.parquet"))) def _prepare_download_parquet( parquet_files, output_dir, count, seed, url_column, text_column, hash_column, ): os.makedirs(output_dir, exist_ok=True) if len(parquet_files) == 1: return parquet_files[0] rng = random.Random(seed) columns = [ url_column, text_column, hash_column, "WIDTH", "HEIGHT", "aesthetic", "punsafe", "pwatermark", ] sampled = _reservoir_sample_parquet( parquet_files=parquet_files, target_count=count, rng=rng, columns=columns, ) if not sampled: raise ValueError("Failed to sample rows from parquet files.") table = pa.Table.from_pylist(sampled) out_path = os.path.join(output_dir, "laion_download_sample.parquet") pq.write_table(table, out_path) logger.info("Wrote sampled parquet list to %s", out_path) return out_path def _detect_image_extension(image): fmt = (image.format or "").upper() if fmt == "JPEG": return "jpg" if fmt == "PNG": return "png" if fmt == "WEBP": return "webp" return "jpg" def _collect_existing_images(image_root, image_extensions): if not image_root or not os.path.exists(image_root): return {} image_map = {} for root, _, files in os.walk(image_root): for name in files: ext = os.path.splitext(name)[1].lower() if ext in image_extensions: stem = os.path.splitext(name)[0] image_map[stem] = os.path.join(root, name) return image_map def _save_image_bytes(image_bytes, output_path): try: with Image.open(BytesIO(image_bytes)) as img: ext = _detect_image_extension(img) if ext == "jpg": img = img.convert("RGB") elif img.mode not in ("RGB", "RGBA"): img = img.convert("RGBA") output_path = os.path.splitext(output_path)[0] + f".{ext}" img.save(output_path) return output_path, img.size except Exception as exc: logger.warning("Failed to decode image bytes: %s", exc) return None, None def _iter_embedded_rows( parquet_files, id_column, caption_column, image_columns, image_bytes_key, image_path_key, ): columns = [id_column, caption_column] + list(image_columns) for parquet_path in parquet_files: parquet_file = pq.ParquetFile(parquet_path) for batch in parquet_file.iter_batches(columns=columns, batch_size=256): batch_dict = batch.to_pydict() batch_len = len(batch) for i in range(batch_len): image_bytes = None image_path = None for image_column in image_columns: image_struct = batch_dict.get(image_column, [None])[i] or {} image_bytes = image_struct.get(image_bytes_key) image_path = image_struct.get(image_path_key) if image_bytes: break if not image_bytes: continue yield { "id": batch_dict.get(id_column, [None])[i], "caption": batch_dict.get(caption_column, [None])[i], "bytes": image_bytes, "path": image_path, } def download_embedded_images( parquet_files, image_root, output_dir, count, seed, sequential, id_column, caption_column, image_columns, image_bytes_key, image_path_key, image_extensions, skip_existing, progress_interval, ): os.makedirs(image_root, exist_ok=True) rng = random.Random(seed) selected_ids = None if not sequential: sampled = _reservoir_sample_parquet( parquet_files=parquet_files, target_count=count, rng=rng, columns=[id_column], ) selected_ids = { str(row.get(id_column)) for row in sampled if row.get(id_column) is not None } if not selected_ids: raise ValueError("Failed to sample IDs from parquet files.") image_extensions = image_extensions or [".jpg", ".png", ".jpeg", ".webp"] existing_map = _collect_existing_images(image_root, image_extensions) if skip_existing else {} if existing_map and len(existing_map) >= count: logger.info( "Found %d existing images in %s (target=%d).", len(existing_map), image_root, count, ) metadata_rows = [] for row in _iter_embedded_rows( parquet_files=parquet_files, id_column=id_column, caption_column=caption_column, image_columns=image_columns, image_bytes_key=image_bytes_key, image_path_key=image_path_key, ): image_id = row.get("id") if image_id is None: continue image_id = str(image_id) if selected_ids is not None and image_id not in selected_ids: continue saved_path = None size = None if image_id in existing_map: saved_path = existing_map[image_id] size = _get_image_size(saved_path) if saved_path is None: shard_dir = image_id[:5] if len(image_id) >= 5 else image_id target_dir = os.path.join(image_root, shard_dir) os.makedirs(target_dir, exist_ok=True) target_path = os.path.join(target_dir, image_id) saved_path, size = _save_image_bytes(row["bytes"], target_path) if not saved_path: continue width, height = size if size else (None, None) metadata_rows.append( { "key": image_id, "caption": row.get("caption"), "status": "success", "width": width, "height": height, } ) if progress_interval and len(metadata_rows) % progress_interval == 0: logger.info("Extracted %d/%d images...", len(metadata_rows), count) if sequential and len(metadata_rows) >= count: break if selected_ids is not None and len(metadata_rows) >= len(selected_ids): break if not metadata_rows: raise ValueError("No embedded images were extracted.") meta_table = pa.Table.from_pylist(metadata_rows) meta_path = os.path.join(image_root, "embedded_metadata.parquet") pq.write_table(meta_table, meta_path) logger.info("Wrote embedded metadata to %s", meta_path) def _reservoir_sample_parquet(parquet_files, target_count, rng, columns): sample = [] total_seen = 0 for parquet_path in parquet_files: parquet_file = pq.ParquetFile(parquet_path) for batch in parquet_file.iter_batches(columns=columns, batch_size=4096): batch_dict = batch.to_pydict() batch_len = len(batch) for i in range(batch_len): row = {col: batch_dict.get(col, [None])[i] for col in columns} total_seen += 1 if len(sample) < target_count: sample.append(row) else: j = rng.randint(0, total_seen - 1) if j < target_count: sample[j] = row return sample def _iter_img2dataset_rows(image_root): parquet_files = sorted(glob.glob(os.path.join(image_root, "*.parquet"))) if not parquet_files: return columns = ["key", "caption", "status", "width", "height"] for parquet_path in parquet_files: parquet_file = pq.ParquetFile(parquet_path) for batch in parquet_file.iter_batches(columns=columns, batch_size=4096): batch_dict = batch.to_pydict() batch_len = len(batch) for i in range(batch_len): status = batch_dict.get("status", [None])[i] if status and status != "success": continue key = batch_dict.get("key", [None])[i] caption = batch_dict.get("caption", [None])[i] width = batch_dict.get("width", [None])[i] height = batch_dict.get("height", [None])[i] if key is None: continue key_str = str(key) yield { "id": key_str, "caption": caption, "width": width, "height": height, } def _image_path_from_id(image_root, key_str, image_extensions): if not key_str: return None shard_dir = key_str[:5] for ext in image_extensions: path = os.path.join(image_root, shard_dir, f"{key_str}{ext}") if os.path.exists(path): return path return os.path.join(image_root, shard_dir, f"{key_str}.jpg") def _round_up_multiple(value, multiple): if multiple <= 1: return int(value) return int(((value + multiple - 1) // multiple) * multiple) def _get_image_size(path): try: with Image.open(path) as img: return img.size except Exception as exc: logger.warning("Failed to read image size for %s: %s", path, exc) return None def build_splits_from_img2dataset( image_root, output_dir, train_count, val_count, test_count, seed, sequential=False, allow_partial=False, id_as_path=False, image_extensions=None, size_multiple=8, ): os.makedirs(output_dir, exist_ok=True) total_needed = train_count + val_count + test_count image_extensions = image_extensions or [".jpg", ".png", ".jpeg", ".webp"] items = [] if sequential: for row in _iter_img2dataset_rows(image_root): items.append(row) if len(items) >= total_needed: break else: rng = random.Random(seed) total_seen = 0 for row in _iter_img2dataset_rows(image_root): total_seen += 1 if len(items) < total_needed: items.append(row) else: j = rng.randint(0, total_seen - 1) if j < total_needed: items[j] = row rng.shuffle(items) if len(items) < total_needed: if not allow_partial: raise ValueError( f"Only found {len(items)} matching images (needed {total_needed})." ) logger.warning( "Only found %d matching images (needed %d).", len(items), total_needed, ) if id_as_path: for item in items: item["id"] = _image_path_from_id(image_root, item["id"], image_extensions) train_items = items[:train_count] val_items = items[train_count : train_count + val_count] test_items = items[train_count + val_count : train_count + val_count + test_count] def write_jsonl(path, rows): with open(path, "w", encoding="utf-8") as f: for row in rows: image_path = row.get("path") if not image_path: image_id = row.get("id") if image_id: if os.path.isabs(image_id): image_path = image_id else: image_path = _image_path_from_id( image_root, image_id, image_extensions ) if image_path: row["path"] = image_path size = _get_image_size(image_path) if size: width, height = size else: width = row.get("width") height = row.get("height") if width and height: row["width"] = _round_up_multiple(int(width), size_multiple) row["height"] = _round_up_multiple(int(height), size_multiple) f.write(json.dumps(row, ensure_ascii=False) + "\n") write_jsonl(os.path.join(output_dir, "train.jsonl"), train_items) write_jsonl(os.path.join(output_dir, "val.jsonl"), val_items) write_jsonl(os.path.join(output_dir, "test.jsonl"), test_items) logger.info( "Wrote splits to %s (train=%d, val=%d, test=%d)", output_dir, len(train_items), len(val_items), len(test_items), ) def _scan_images(image_root, image_extensions): if not image_root or not os.path.exists(image_root): return {} image_map = {} for root, _, files in os.walk(image_root): for name in files: ext = os.path.splitext(name)[1].lower() if ext in image_extensions: stem = os.path.splitext(name)[0] image_map[stem] = os.path.join(root, name) return image_map def _collect_metadata( parquet_files, image_map, target_count, url_column, text_column, hash_column, ): selected = [] hashes = set(image_map.keys()) if not hashes: return selected columns = [ hash_column, url_column, text_column, "WIDTH", "HEIGHT", "aesthetic", "punsafe", "pwatermark", ] for parquet_path in parquet_files: parquet_file = pq.ParquetFile(parquet_path) for batch in parquet_file.iter_batches(columns=columns, batch_size=4096): batch_dict = batch.to_pydict() for i in range(len(batch)): hash_value = batch_dict.get(hash_column, [None])[i] if hash_value is None: continue hash_str = str(hash_value) path = image_map.get(hash_str) if not path: continue selected.append( { "file": path, "url": batch_dict.get(url_column, [None])[i], "text": batch_dict.get(text_column, [None])[i], "width": batch_dict.get("WIDTH", [None])[i], "height": batch_dict.get("HEIGHT", [None])[i], "hash": hash_str, "aesthetic": batch_dict.get("aesthetic", [None])[i], "punsafe": batch_dict.get("punsafe", [None])[i], "pwatermark": batch_dict.get("pwatermark", [None])[i], } ) if len(selected) >= target_count: return selected return selected def build_splits( data_files, image_root, image_extensions, output_dir, train_count, val_count, test_count, seed, url_column, text_column, hash_column, sequential=False, allow_partial=False, size_multiple=8, ): os.makedirs(output_dir, exist_ok=True) parquet_files = _expand_parquet_files(data_files) if not parquet_files: raise ValueError("No parquet files found. Check --data-files.") image_map = _scan_images(image_root, image_extensions) if not image_map: raise ValueError("No images found in image_root.") total_needed = train_count + val_count + test_count logger.info( "Collecting %d samples from %d parquet files (images=%d)", total_needed, len(parquet_files), len(image_map), ) items = _collect_metadata( parquet_files=parquet_files, image_map=image_map, target_count=total_needed, url_column=url_column, text_column=text_column, hash_column=hash_column, ) if len(items) < total_needed: if not allow_partial: raise ValueError( f"Only found {len(items)} matching images (needed {total_needed})." ) logger.warning( "Only found %d matching images (needed %d).", len(items), total_needed, ) if not sequential: rng = random.Random(seed) rng.shuffle(items) train_items = items[:train_count] val_items = items[train_count : train_count + val_count] test_items = items[train_count + val_count : train_count + val_count + test_count] def write_jsonl(path, rows): with open(path, "w", encoding="utf-8") as f: for row in rows: image_path = row.get("path") or row.get("file") if image_path: row["path"] = image_path size = _get_image_size(image_path) if size: width, height = size else: width = row.get("width") height = row.get("height") if width and height: row["width"] = _round_up_multiple(int(width), size_multiple) row["height"] = _round_up_multiple(int(height), size_multiple) f.write(json.dumps(row, ensure_ascii=False) + "\n") write_jsonl(os.path.join(output_dir, "train.jsonl"), train_items) write_jsonl(os.path.join(output_dir, "val.jsonl"), val_items) write_jsonl(os.path.join(output_dir, "test.jsonl"), test_items) logger.info( "Wrote splits to %s (train=%d, val=%d, test=%d)", output_dir, len(train_items), len(val_items), len(test_items), ) if __name__ == "__main__": main() ''' python -m tools.sample_backgrounds \ --download \ --count 20100 \ --build-splits \ --train-count 19000 \ --val-count 1000 \ --test-count 200 \ --data-files "/project/llmsvgen/share/data/kmw_layered_dataset/laion2B-en-aesthetic-image/*.parquet" \ --image-root "/project/llmsvgen/share/data/kmw_layered_dataset/laion2B-en-aesthetic-image" \ --output-dir "/project/llmsvgen/jinmin/SynLayers/data/laion2b_splits" python -m tools.sample_backgrounds \ --download \ --build-splits \ --count 40200 \ --sequential \ --id-as-path \ --train-count 19000 \ --val-count 1000 \ --test-count 200 \ --data-files "/project/llmsvgen/share/data/kmw_layered_dataset/PrismLayersPro-image/data/*.parquet" \ --image-root "/project/llmsvgen/share/data/kmw_layered_dataset/PrismLayersPro-image/data/haolin/PrismLayersPro-image" \ --output-dir "/project/llmsvgen/jinmin/SynLayers/data/prismlayerspro_splits" '''