File size: 4,197 Bytes
69def8e | 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 | import argparse
import os
import sys
import yaml
import numpy as np
# Python 3.12 removed distutils, but tensorflow still requires it.
try:
import distutils
except ImportError:
import setuptools
sys.modules['distutils'] = setuptools._distutils
import tensorflow_datasets as tfds
from PIL import Image
def load_config(config_path):
with open(config_path, 'r') as f:
return yaml.safe_load(f)
def set_seed(seed):
np.random.seed(seed)
def main():
parser = argparse.ArgumentParser(description="Ingest LFW dataset deterministically.")
parser.add_argument("--config", type=str, required=True, help="Path to the config file.")
args = parser.parse_args()
config = load_config(args.config)
seed = config.get("seed", 42)
set_seed(seed)
print(f"Loading {config['data_source']}...")
# Load dataset. LFW provides 'train' split default usually, but tfds lfw might not have pre-defined splits other than 'train'
ds, ds_info = tfds.load("lfw", split="train", with_info=True, as_supervised=False)
# We will build a dictionary of identity -> list of (image_array, original_filename)
# LFW in TFDS has 'label' which is integer, and 'image' which is array.
# We need the actual string labels to sort them properly if we want identity names,
# or we can just use the integer label if we sort deterministically.
# LFW feature 'label' contains the text identity
dataset_dict = {}
print("Iterating over dataset to group by identity...")
for i, example in enumerate(tfds.as_numpy(ds)):
identity = example['label'].decode('utf-8')
image = example['image']
if identity not in dataset_dict:
dataset_dict[identity] = []
dataset_dict[identity].append(image)
identities = sorted(list(dataset_dict.keys()))
print(f"Found {len(identities)} identities and {sum(len(v) for v in dataset_dict.values())} total images.")
# Shuffle identities deterministically
np.random.shuffle(identities)
# Simple split of identities
n_total = len(identities)
n_train = int(n_total * config["train_ratio"])
n_val = int(n_total * config["val_ratio"])
train_ids = identities[:n_train]
val_ids = identities[n_train:n_train + n_val]
test_ids = identities[n_train + n_val:]
splits = {
"train": train_ids,
"val": val_ids,
"test": test_ids
}
output_dir = config["output_dir"]
os.makedirs(output_dir, exist_ok=True)
counts = {}
print("Saving images to disk...")
for split_name, ids in splits.items():
split_dir = os.path.join(output_dir, split_name)
os.makedirs(split_dir, exist_ok=True)
num_images = 0
for identity in ids:
identity_dir = os.path.join(split_dir, identity)
os.makedirs(identity_dir, exist_ok=True)
images = dataset_dict[identity]
# Since tfds doesn't guarantee order inside identity, we ensure deterministic order by... wait, tfds arrays are indistinguishable.
# We'll just save them in the order they appeared, which is deterministic given TFDS yields deterministically.
for i, img_arr in enumerate(images):
img = Image.fromarray(img_arr)
img_path = os.path.join(identity_dir, f"{identity}_{i:04d}.jpg")
img.save(img_path)
num_images += 1
counts[split_name] = {
"identities": len(ids),
"images": num_images
}
print(f"{split_name}: {len(ids)} identities, {num_images} images.")
manifest = {
"seed": seed,
"split_policy": config["split_policy"],
"counts": counts,
"data_source": config["data_source"],
"cache_location": output_dir
}
manifest_path = config["manifest_path"]
os.makedirs(os.path.dirname(manifest_path), exist_ok=True)
with open(manifest_path, 'w') as f:
yaml.dump(manifest, f, sort_keys=False)
print(f"Manifest written to {manifest_path}")
if __name__ == "__main__":
main()
|