Search code examples
object-detection

Object Detection MASK RCNN only for 2 classes


I use (https://github.com/matterport/Mask_RCNN) MASKRCNN for object detection and everything works fine. Is there a way to detect only certain objects? There are more than 80 classes of which I only need 2 pieces (like car and person). I would like to have the remaining classes not detected. How can they be removed?


Solution

  • I guess you are using the default pre-trained coco model for detection which comes with 80 classes. You can train your own model. First of all you have to use VIA (VGG image annotator) to label the classes you want to predict. Once that is done, you have to make some changes to the code of the model. For example, if you are using the file "balloon.py", you have to add classes, make some changes in load_mask() function and few other parts of the codes. After that you can start training your model, then use for detection and segmentation.

    Yes, and for VIA, try using the version 1.0.0 as the format of the .json file is slightly changed in in the updated versions, which generally makes them incompatible with training on your custom data sets.

    Check one example here