Маска R-CNN, Мультикласс

Question

Я следую за классным c примером воздушного шара Маски R-CNN. Я тренировал модель с 2 классами, транспортным средством и bottle. Теперь модель прекрасно обнаруживает оба объекта, но для каждого из них указывается транспортное средство (class_id = 1), где все бутылки называются транспортными. Я, вероятно, допустил некоторые ошибки при настройке кода. Кто-нибудь имеет представление об этой проблеме?

Изменено количество классов в config.py (1 + nb классов). Добавлены эти классы в load_ в классе пользовательских наборов данных

Но я думаю, Мне нужно сделать еще одну вещь, и это примерно здесь:

 for a in annotations:
        # print(a)
        # Get the x, y coordinates of points of the polygons that make up
        # the outline of each object instance. There are stores in the
        # shape_attributes (see json format above)
        polygons = [r['shape_attributes'] for r in a['regions'].values()]

        # load_mask() needs the image size to convert polygons to masks.
        # Unfortunately, VIA doesn't include it in JSON, so we must read
        # the image. This is only managable since the dataset is tiny.
        image_path = os.path.join(dataset_dir, a['filename'])
        image = skimage.io.imread(image_path)
        height, width = image.shape[:2]

        self.add_image(
            "vehicle",  ## for a single class just add the name here
            image_id=a['filename'],  # use file name as a unique image id
            path=image_path,
            width=width, height=height,
            polygons=polygons)

Я не мог понять, что изменить здесь. Спасибо за ваше время.

Answer 1

Каждый раз, когда вы вызываете функцию self.add_image(....) в вашем for loop, вы должны проверить, к какому классу принадлежит ваше изображение. Если он принадлежит vehicle, вы звоните:

self.add_image("vehicle", ...)

Если он принадлежит bottle, вам следует позвонить:

self.add_image("bottle", ...)

Кроме того, в вашем файле config, в дополнение к установке NUM_CLASSES на 1+nb_classes, вы должны дать вашей конфигурации имя, например, "vehicle_and_bottle". Затем в вашей функции load вы должны добавить эти:

self.add_class("vehicle_and_bottle", 1, "bottle") # first class
self.add_class("vehicle_and_bottle", 2, "vehicle") # second class