diff --git a/api/src/yolo/merge.py b/api/src/yolo/merge.py new file mode 100644 index 0000000..12bfafc --- /dev/null +++ b/api/src/yolo/merge.py @@ -0,0 +1,94 @@ +"""Combine COCO folders (images + ) 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 "__" 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)}")