Маска RCNN, как изменить код load_mask в примере всплывающей подсказки для 2 классов

Question

Я использую маску RCNN и пытаюсь изменить этот пример (https://github.com/matterport/Mask_RCNN/blob/master/samples/balloon/balloon.py), который идентифицирует воздушный шар, чтобы он идентифицировал воздушный шар и торт, т.е. увеличил количество классов до двух.

def load_mask(self, image_id):
    """Generate instance masks for an image.
   Returns:
    masks: A bool array of shape [height, width, instance count] with
        one mask per instance.
    class_ids: a 1D array of class IDs of the instance masks.
    """
    # If not a balloon dataset image, delegate to parent class.
    image_info = self.image_info[image_id]
    if image_info["source"] != "student":
        return super(self.__class__, self).load_mask(image_id)

    # Convert polygons to a bitmap mask of shape
    # [height, width, instance_count]
    info = self.image_info[image_id]
    mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
                    dtype=np.uint8)
    for i, p in enumerate(info["polygons"]):
        # Get indexes of pixels inside the polygon and set them to 1
        rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])
        mask[rr, cc, i] = 1

    **# Return mask, and array of class IDs of each instance. Since we have
    # one class ID only, we return an array of 1s**
    return mask.astype(np.bool), np.ones([mask.shape[-1]], dtype=np.int32)

Как видно из последнего комментария в функции load_mask, этот код написан только для одного класса. Как мне изменить его для двух классов?

Answer 1

, внимательно посмотрите на эту проблему: https://github.com/matterport/Mask_RCNN/issues/372 Моя выглядит так:

def load_mask(self, image_id):
    """Generate instance masks for an image.
   Returns:
    masks: A bool array of shape [height, width, instance count] with
        one mask per instance.
    class_ids: a 1D array of class IDs of the instance masks.
    """
    # If not a balloon dataset image, delegate to parent class.
    info = self.image_info[image_id]
    if info["source"] != "sun":
        return super(self.__class__, self).load_mask(image_id)
    class_ids = info['class_ids']

    # Convert polygons to a bitmap mask of shape
    # [height, width, instance_count]
    mask = np.zeros([info["height"], info["width"], len(info["polygons"])],
                    dtype=np.uint8)
    for i, p in enumerate(info["polygons"]):
        # Get indexes of pixels inside the polygon and set them to 1
        rr, cc = skimage.draw.polygon(p['all_points_y'], p['all_points_x'])
        mask[rr, cc, i] = 1

    # Return mask, and array of class IDs of each instance. Since we have
    # one class ID only, we return an array of 1s
    class_ids = np.array(class_ids, dtype=np.int32)
    return mask.astype(np.bool), class_ids

Не забудьте изменить Load_ballon:

def load_balloon(self, dataset_dir, subset):
    """Load a subset of the Balloon dataset.
    dataset_dir: Root directory of the dataset.
    subset: Subset to load: train or val
    """
    # Add classes. We have only one class to add.
    self.add_class("balloon", 1, "ballon")
    self.add_class("balloon", 2, "cakes")



    # Train or validation dataset?
    assert subset in ["train", "val"]
    dataset_dir = os.path.join(dataset_dir, subset)

    # Load annotations
    #...........

    # We mostly care about the x and y coordinates of each region
    annotations = json.load(open(os.path.join(dataset_dir, "via_region_data.json")))
    annotations = list(annotations.values())  # don't need the dict keys

    # The VIA tool saves images in the JSON even if they don't have any
    # annotations. Skip unannotated images.
    annotations = [a for a in annotations if a['regions']]

    # Add images
    for a in annotations:
        # Get the x, y coordinaets 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()]
        objects = [s['region_attributes'] for s in a['regions'].values()]
        class_ids = [int(n['class']) for n in objects]
        # 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(
            "balloon",
            image_id=a['filename'],  # use file name as a unique image id
            path=image_path,
            width=width, height=height,
            polygons=polygons,
            class_ids=class_ids)

Answer 2

вы можете использовать версию маски rcnn для фигур в поездах для 2 или более классов, так как теперь набор данных воздушного шара готовится с использованием одного идентификатора класса, т.е. воздушного шара