Spaces:
Sleeping
Sleeping
| """Build a Food-101 subset for training. | |
| After ``scripts/download_data.py --cv`` extracted the 10 target classes into | |
| ``data/raw/food101_subset/food-101/images/<class>/<image>.jpg``, this script | |
| splits them into train/val/test directories using Food-101's own meta files. | |
| Produces: | |
| data/processed/cv/train/<class>/*.jpg | |
| data/processed/cv/val/<class>/*.jpg | |
| data/processed/cv/test/<class>/*.jpg | |
| Usage: | |
| python -m src.cv.prepare_data | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import shutil | |
| import sys | |
| from pathlib import Path | |
| from random import Random | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) | |
| from src.config import CV_CLASSES_PATH, CV_TARGET_CLASSES, PROCESSED_DIR, RAW_DIR # noqa: E402 | |
| CV_RAW = RAW_DIR / "food101_subset" / "food-101" | |
| CV_OUT = PROCESSED_DIR / "cv" | |
| def _read_meta(path: Path) -> list[str]: | |
| return [line.strip() for line in path.read_text().splitlines() if line.strip()] | |
| def _link_image(src: Path, dst: Path) -> None: | |
| dst.parent.mkdir(parents=True, exist_ok=True) | |
| if dst.exists(): | |
| return | |
| try: | |
| dst.symlink_to(src) | |
| except OSError: | |
| shutil.copy2(src, dst) | |
| def main(val_fraction: float = 0.15, seed: int = 42) -> None: | |
| train_meta = CV_RAW / "meta" / "train.txt" | |
| test_meta = CV_RAW / "meta" / "test.txt" | |
| if not train_meta.exists() or not test_meta.exists(): | |
| raise FileNotFoundError( | |
| "Food-101 meta files missing. Run 'python scripts/download_data.py --cv'." | |
| ) | |
| targets = set(CV_TARGET_CLASSES) | |
| rnd = Random(seed) | |
| train_entries = [e for e in _read_meta(train_meta) if e.split("/")[0] in targets] | |
| test_entries = [e for e in _read_meta(test_meta) if e.split("/")[0] in targets] | |
| val_entries: list[str] = [] | |
| new_train: list[str] = [] | |
| by_class: dict[str, list[str]] = {} | |
| for entry in train_entries: | |
| by_class.setdefault(entry.split("/")[0], []).append(entry) | |
| for cls, items in by_class.items(): | |
| rnd.shuffle(items) | |
| cut = int(len(items) * val_fraction) | |
| val_entries.extend(items[:cut]) | |
| new_train.extend(items[cut:]) | |
| splits = {"train": new_train, "val": val_entries, "test": test_entries} | |
| for split, entries in splits.items(): | |
| for entry in entries: | |
| src = CV_RAW / "images" / f"{entry}.jpg" | |
| cls = entry.split("/")[0] | |
| dst = CV_OUT / split / cls / f"{entry.split('/')[1]}.jpg" | |
| if src.exists(): | |
| _link_image(src, dst) | |
| print(f"[cv.prepare_data] {split:>5}: {len(entries):,} images") | |
| CV_CLASSES_PATH.parent.mkdir(parents=True, exist_ok=True) | |
| CV_CLASSES_PATH.write_text(json.dumps(sorted(CV_TARGET_CLASSES), indent=2)) | |
| print(f"[cv.prepare_data] wrote class list to {CV_CLASSES_PATH}") | |
| if __name__ == "__main__": | |
| main() | |