merge

api/src/yolo/merge.py

@@ -0,0 +1,94 @@
+"""Combine COCO folders (images + <annotation_file>) into one COCO folder
+that YoloDataset reads, e.g. WIDER and faces_scenery into one training set.
+
+- Images are hard-linked into the output folder (copied if the folders are
+  on different drives), named "<source>__<file_name>" so names never clash.
+- Image and annotation ids are renumbered.
+- Categories are merged by name and numbered 1, 2, ... in order of first
+  appearance, so the same class in both sources becomes one class. --labels
+  must equal the number of merged categories.
+"""
+
+import json
+import os
+import re
+import shutil
+from pathlib import Path
+
+from src.yolo.config import Config
+
+
+def source_names(folders):
+    """Readable, unique prefix per source: data/wider_faces_coco/train -> wider_faces_coco_train."""
+    names = []
+    for folder in folders:
+        name = re.sub(r"[^\w.-]+", "_", "_".join(folder.parts[-2:]))
+        while name in names:
+            name += "_"
+        names.append(name)
+    return names
+
+
+def link_or_copy(source: Path, target: Path):
+    try:
+        os.link(source, target)
+    except OSError:
+        shutil.copy2(source, target)
+
+
+def merge_coco(cfg: Config):
+    """Combine the COCO folders in --merge-dirs (comma separated) into --merge-out-dir."""
+    folders = [Path(d.strip()) for d in cfg.merge_dirs.split(",") if d.strip()]
+    if len(folders) < 2:
+        raise SystemExit("Set --merge-dirs to at least two COCO folders, e.g. data/a/train,data/b/train")
+    if not cfg.merge_out_dir:
+        raise SystemExit("Set --merge-out-dir to the folder the combined dataset is written to")
+    out = Path(cfg.merge_out_dir)
+    if out.exists():
+        raise SystemExit(f"{out} already exists, delete it or pick another --merge-out-dir")
+
+    # read and check everything before writing anything
+    sources = []
+    for folder in folders:
+        annotation_path = folder / cfg.annotation_file
+        if not annotation_path.is_file():
+            raise SystemExit(f"{annotation_path} not found")
+        coco = json.loads(annotation_path.read_text())
+        missing = [i["file_name"] for i in coco["images"] if not (folder / i["file_name"]).is_file()]
+        if missing:
+            raise SystemExit(f"{folder}: {len(missing)} images listed in {cfg.annotation_file} "
+                             f"are missing, e.g. {missing[0]}")
+        sources.append(coco)
+
+    categories = {}  # name -> merged category
+    images, annotations = [], []
+    out.mkdir(parents=True)
+
+    for folder, name, coco in zip(folders, source_names(folders), sources):
+        category_ids = {}
+        for category in coco["categories"]:
+            merged = categories.setdefault(category["name"], {**category, "id": len(categories) + 1})
+            category_ids[category["id"]] = merged["id"]
+
+        image_ids = {}
+        for image in coco["images"]:
+            # flat output folder: keep subfolders of the source in the name
+            file_name = f"{name}__{Path(image['file_name']).as_posix().replace('/', '__')}"
+            link_or_copy(folder / image["file_name"], out / file_name)
+            image_ids[image["id"]] = len(images) + 1
+            images.append({**image, "id": len(images) + 1, "file_name": file_name})
+
+        for annotation in coco["annotations"]:
+            annotations.append({**annotation, "id": len(annotations) + 1,
+                                "image_id": image_ids[annotation["image_id"]],
+                                "category_id": category_ids[annotation["category_id"]]})
+
+        print(f"{folder}: {len(coco['images'])} images, {len(coco['annotations'])} annotations "
+              f"(prefix {name}__)")
+
+    merged = {"images": images, "annotations": annotations, "categories": list(categories.values())}
+    (out / cfg.annotation_file).write_text(json.dumps(merged))
+
+    names = ", ".join(f"{c['id']} = {c['name']}" for c in merged["categories"])
+    print(f"{out}: {len(images)} images, {len(annotations)} annotations, categories: {names}")
+    print(f"Train with --labels {len(categories)}")