Generic Loader Function for MS COCO Style Dataset

I am working with Mask-RCNN and want to train my own coco-style dataset with few categories. For the start, I have only 2 classes (apart from the background).

Although Mask-RCNN comes with sample dataset, they either contain only one class or they generate the data themselves, making it impossible for me to understand how to load the annotated images and start the training. I've been spending a lot of time in Python functions but kept getting errors such as:

TypeError: list indices must be integers or slices, not str

which left me pretty much clueless.

The current loader I have (to load my coco-style dataset) is the following (from coco loader in Mask R-CNN repository)

 def load_components(self, dataset_dir, subset, year=DEFAULT_DATASET_YEAR):

    coco = COCO("{}/annotations/instances_{}{}.json".format(dataset_dir, subset, year))
    if subset == "minival" or subset == "valminusminival":
        subset = "val"
    image_dir = "{}/{}{}".format(dataset_dir, subset, year)

    # Load all classes or a subset?
    class_ids = sorted(coco.getCatIds())

    # All images or a subset?
    image_ids = list(coco.imgs.keys())

    # Add classes
    for i in class_ids:
        self.add_class("coco", i, coco.loadCats(i)[0]["name"])

    # Add images
    for i in image_ids:
        self.add_image(
            "coco", image_id=i,
            path=os.path.join(image_dir, coco.imgs[i]['file_name']),
            width=coco.imgs[i]["width"],
            height=coco.imgs[i]["height"],
            annotations=coco.loadAnns(coco.getAnnIds(
                imgIds=[i], catIds=class_ids, iscrowd=None)))

Which does not load anything, naturally.

Can someone help me write a loader function which can be used in all kinds of coco-style datasets, regardless of class number?

Here is my dataset structure, which is coco-style:

components
│
└───train
│   │ 
│   └───annotations
│   │    │ <image_id>_<object_class_name>_<annotation_id>.png
│   │    │ ...
│   │   
│   └───<subset><year>
│   │    │   <image_id>.jpeg
│   │    │   ...  
│   └───  instances_components_train2018.json
└───val
    │ 
    └───annotations
    │    │ <image_id>_<object_class_name>_<annotation_id>.png
    │    │ ...
    │   
    └───<subset><year>
    │    │   <image_id>.jpeg
    │    │   ...
    │    │  
    └───  instances_components_val2018.json

EDIT: Full traceback:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-11-5f1fb6581c25> in <module>()
      1 # Load dataset
      2 dataset = components.ComponentsDataset()
----> 3 dataset.load_components(DATASET_DIR, subset="train")
      4 
      5 # Must call before using the dataset

~/workspace/Mask_RCNN/samples/components/components.py in load_components(self, dataset_dir, subset)
     72         # Load all classes or a subset?
     73         class_ids = sorted(coco.getCatIds())
---> 74 
     75         # All images or a subset?
     76         image_ids = list(coco.imgs.keys())

~/workspace/Mask_RCNN/samples/components/components.py in <listcomp>(.0)
     72         # Load all classes or a subset?
     73         class_ids = sorted(coco.getCatIds())
---> 74 
     75         # All images or a subset?
     76         image_ids = list(coco.imgs.keys())

TypeError: list indices must be integers or slices, not str

Okay so I figured it out. There is a file which I found here, showing a generic way of loading a coco-style dataset and making it work.

Of course, if you want to do this, you need to modify the variables a bit, since originally it was designed for "shapes" dataset. So, you rename everything that has "shapes" to whatever your dataset name is.

DO NOT attempt to change the functions and way they work. Do not try to copy something else into that code, the whole structure is a chain one and if you change something in the function (in terms of flow) you may not be able to get the code running.