Eave_Z
Eave_Z

Reputation: 5

add a new class on pre trained mask rcnn

I want fine tuned mrcnn on my custom dataset, and after tuning, it could only detect the class I add, and my tuning based on MS coco weight, here is my code:

class ShoeConfig(Config):
    NAME = "shoe"

    # We use a GPU with 12GB memory, which can fit two images.
    # Adjust down if you use a smaller GPU.
    IMAGES_PER_GPU = 2

    # Number of classes (including background)
    NUM_CLASSES = 1 + 1  # Background + shoe

    # Number of training steps per epoch
    STEPS_PER_EPOCH = 11797

    # Skip detections with < 90% confidence
    DETECTION_MIN_CONFIDENCE = 0.9

which is config part,

class ShoeDataset(utils.Dataset):

    def load_shoe(self, dataset_dir, subset):
        
        # Add classes. We have only one class to add.
        self.add_class("shoe", 1, "shoe")

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

which is load dataset code snippet, are there errors, and how should I modify it?

Upvotes: 0

Views: 89

Answers (0)

Related Questions