make forward train signature compatible with single stage detectors
Summary: gt_masks isn't an arg for single stage detectors so if no gt_masks are available we don't pass it as an arg.
@@ -199,8 +199,6 @@ class MMDetDetector(nn.Module):
return mm_PolygonMasks(m.polygons, shape[0], shape[1])
gt_masks = [convert_mask(x.gt_masks, x.image_size) for x in gt_instances]
- else:
- gt_masks = None
losses_and_metrics = self.detector.forward_train(
images,
metas,
@@ -208,6 +206,13 @@ class MMDetDetector(nn.Module):
[x.gt_classes for x in gt_instances],
gt_masks=gt_masks,
)
+ else:
+ losses_and_metrics = self.detector.forward_train(
+ images,
+ metas,
+ [x.gt_boxes.tensor for x in gt_instances],
+ [x.gt_classes for x in gt_instances],
+ )
return _parse_losses(losses_and_metrics)
else:
results = self.detector.simple_test(images, metas, rescale=rescale)
