Search code examples
computer-visionconv-neural-networkimage-segmentationfaster-rcnndetectron

Train Mask R-CNN with some classes having no masks but only bounding boxes

I want to train a model to detect three different types of defects. I have a training dataset where two of these classes contain segmentation masks, but one contains only bounding boxes. Can I train a shared model or do I need to separate the training dataset and train a Faster R-CNN and a Mask R-CNN?

(I only care about bounding box output for the class containing no masks in the training data.)

Solution

  • You can create 'weak' masks from those bounding boxes and then combine those two datasets. Something like below:

    mask = np.zeros((256, 256), dtype=np.float32)
mask[y:y+h, x:x+w] = 255.

    If the two datasets are small, combining them will yield better results. But if the datasets are big enough (>2000 images) then you can use FasterRCNN + MaskRCNN approach.