Hi,

Awesome work. Something similar to what I did, though for semantic segmentation.

A small suggestion: From what I understand, you have hard-coded the Fast-RCNN predictor.

  self._model.roi_heads.box_predictor = FastRCNNPredictor(in_features, len(classes) + 1)

Having the ability to plugin users' own models would be better, I believe. This would translate to making the model class more abstract.

Well done!