mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-07-23 02:02:08 +08:00
firestar5683
81 lines
2.9 KiB
Python
81 lines
2.9 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import shutil
|
|
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
|
|
if __package__ in (None, ""):
|
|
import sys
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from common import EXTENDED_CLASSIFIER_SPEED_VALUES # type: ignore # noqa: TID251
|
|
else:
|
|
from .common import EXTENDED_CLASSIFIER_SPEED_VALUES
|
|
|
|
|
|
IMAGE_SUFFIXES = frozenset((".jpg", ".jpeg", ".png", ".bmp", ".webp"))
|
|
SUPPORTED_CLASSES = frozenset(str(value) for value in EXTENDED_CLASSIFIER_SPEED_VALUES) | {"reject"}
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Merge image-folder classifier datasets using hard links when possible.")
|
|
parser.add_argument("sources", nargs="+", type=Path, help="Classifier roots containing train/ and val/.")
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
parser.add_argument("--overwrite", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
def link_or_copy(source: Path, destination: Path) -> None:
|
|
destination.parent.mkdir(parents=True, exist_ok=True)
|
|
try:
|
|
os.link(source, destination)
|
|
except OSError:
|
|
# copy2 preserves macOS metadata as AppleDouble `._` files on exFAT. Image
|
|
# loaders then mistake those sidecars for corrupt training images.
|
|
shutil.copyfile(source, destination)
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
sources = [path.expanduser().resolve() for path in args.sources]
|
|
output = args.output.expanduser().resolve()
|
|
if output.exists():
|
|
if not args.overwrite:
|
|
raise FileExistsError(f"Output already exists: {output}")
|
|
shutil.rmtree(output)
|
|
|
|
counts: Counter[str] = Counter()
|
|
for source_index, source_root in enumerate(sources):
|
|
for split in ("train", "val"):
|
|
split_root = source_root / split
|
|
if not split_root.is_dir():
|
|
raise FileNotFoundError(split_root)
|
|
for class_dir in sorted(split_root.iterdir()):
|
|
if not class_dir.is_dir() or class_dir.name.startswith("._"):
|
|
continue
|
|
if class_dir.name not in SUPPORTED_CLASSES:
|
|
raise ValueError(f"Unsupported classifier class {class_dir.name!r} in {source_root}")
|
|
for image_path in sorted(class_dir.iterdir()):
|
|
if image_path.name.startswith("._") or image_path.suffix.lower() not in IMAGE_SUFFIXES or not image_path.is_file():
|
|
continue
|
|
destination = output / split / class_dir.name / f"s{source_index:02d}_{image_path.name}"
|
|
link_or_copy(image_path, destination)
|
|
counts[f"{split}/{class_dir.name}"] += 1
|
|
|
|
summary = {
|
|
"sources": [str(path) for path in sources],
|
|
"output": str(output),
|
|
"counts": dict(sorted(counts.items())),
|
|
}
|
|
(output / "merge_summary.json").write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="ascii")
|
|
print(json.dumps(summary, indent=2, sort_keys=True))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|