Как обучить пользовательский детектор ключевых точек для оценки позы дрона. Detectron2

Поскольку я не мог найти ответа в другом месте, я решил описать свою проблему здесь. Я пытаюсь создать детектор ключевых точек Drone Eachine TrashCan для оценки его позы. Я следил за некоторыми уроками. Первый был с TensorFlow ObjectDetectionAPI, и поскольку я не смог найти с его помощью решения, я попытался использовать detectron2. Все было хорошо, пока мне не понадобилось зарегистрировать собственный набор данных для переобучения модели.

Я запускаю код в Google Colab и использую coco-annotator для создания аннотаций (https://github.com/jsbroks/coco-annotator/)

Я не думаю, что я неправильно аннотировал свой набор данных, но кто знает, я покажу в гиперссылке ниже, чтобы немного представить ее вам: Картинка с аннотациями, сделанными мной

Я использовал этот код для регистрации данных:

from detectron2.data.datasets import register_coco_instances
register_coco_instances("TrashCan_train", {}, "./TrashCan_train/mask_train.json", "./TrashCan_train")
register_coco_instances("TrashCan_test", {}, "./TrashCan_test/mask_test.json", "./TrashCan_test")

Это не выдает мне ошибки, но когда я пытаюсь начать процесс обучения с этим кодом:

from detectron2.engine import DefaultTrainer

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x.yaml"))
cfg.DATASETS.TRAIN = ("TrashCan_train",)
cfg.DATASETS.TEST = ("TrashCan_test")
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = ("detectron2://COCO-Keypoints/keypoint_rcnn_R_50_FPN_3x/137849621/model_final_a6e10b.pkl")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025  # pick a good LR
cfg.SOLVER.MAX_ITER = 25    # 300 iterations seems good enough for this toy dataset; you may need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128   # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 1  # liczba klas

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg) 
trainer.resume_or_load(resume=False)
trainer.train()

Я получаю вот что:

WARNING [07/14 14:36:52 d2.data.datasets.coco]: 
Category ids in annotations are not in [1, #categories]! We'll apply a mapping for you.

[07/14 14:36:52 d2.data.datasets.coco]: Loaded 5 images in COCO format from ./mask_train/mask_train.json

---------------------------------------------------------------------------

KeyError                                  Traceback (most recent call last)

<ipython-input-12-f4f5153c62a1> in <module>()
     14 
     15 os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
---> 16 trainer = DefaultTrainer(cfg)
     17 trainer.resume_or_load(resume=False)
     18 trainer.train()

7 frames

/usr/local/lib/python3.6/dist-packages/detectron2/data/datasets/coco.py in load_coco_json(json_file, image_root, dataset_name, extra_annotation_keys)
    183             obj["bbox_mode"] = BoxMode.XYWH_ABS
    184             if id_map:
--> 185                 obj["category_id"] = id_map[obj["category_id"]]
    186             objs.append(obj)
    187         record["annotations"] = objs

KeyError: 9

Здесь вы можете скачать мои файлы:

https://github.com/BrunoKryszkiewicz/EachineTrashcan-keypoints/blob/master/TrashCan_Datasets.zip

а где вы найдете мою записную книжку Colab:

https://colab.research.google.com/drive/1AlLZxrR64irms9mm-RiZ4DNOs2CaAEZ9?usp=sharing

Я создал такой небольшой набор данных, потому что в первую очередь я хотел пройти через тренировочный процесс. Если мне это удастся, я увеличу свой набор данных.

Теперь я даже не уверен, можно ли переобучить человеческую модель ключевых точек для получения ключевых точек беспилотного объекта или любого другого. Я должен сказать, что я новичок в этой теме, поэтому прошу вашего понимания. Если вы знаете какие-либо руководства по созданию пользовательского (не человеческого) детектора ключевых точек, я буду благодарен за любую информацию.

С наилучшими пожеланиями!