Add end-to-end polygon (Mask R-CNN) workflow#1399
Conversation
Implement instance segmentation paralleling the box and point workflows: - Add Mask R-CNN model (task="polygon"), encapsulating the torchvision background-class convention by shifting labels in forward() - Refactor PolygonDataset to mirror BoxDataset (csv_file, 3-tuple, joint image/box/mask augmentation) - Support polygon prediction formatting and mask-based geometry detection - Wire segmentation mAP and polygon recall/precision into validation - Enable polygon mosaicking in predict_tile via reduce_polygons - Add tests for the dataset, model, formatting, training and prediction Co-authored-by: Cursor <[email protected]>
Co-authored-by: Cursor <[email protected]>
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1399 +/- ##
==========================================
- Coverage 86.61% 85.71% -0.90%
==========================================
Files 26 28 +2
Lines 3736 3970 +234
==========================================
+ Hits 3236 3403 +167
- Misses 500 567 +67
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
There was a problem hiding this comment.
@bw4sz as far as your concerns about this, my opinion is the general feel of this PR is good. It cribs the right bits from other sections of the code. Can you feed the review notes back into Cursor?
|
|
||
| # Compute precision, recall and empty frame metrics. | ||
| if self.model.task == "box" or self.model.task == "point": | ||
| if self.model.task in ("box", "point", "polygon"): |
There was a problem hiding this comment.
Redundant? At least for now we can compute PR for any task so we can always add this metric.
| """Determine the geometry type of the batched result.""" | ||
| # Assumes that all geometries are the same in a batch | ||
| if "boxes" in batched_result.keys(): | ||
| # Assumes that all geometries are the same in a batch. Mask R-CNN |
There was a problem hiding this comment.
I would remove the reference to Mask R-CNN here and say "Polygon model results that contain both ... "
| label_dict: dict = None, | ||
| **kwargs, | ||
| ): | ||
| backbone = torchvision.models.detection.maskrcnn_resnet50_fpn( |
There was a problem hiding this comment.
There is also V2. https://docs.pytorch.org/vision/main/models/mask_rcnn.html
Improved Mask R-CNN model with a ResNet-50-FPN backbone from the Benchmarking Detection Transfer Learning with Vision Transformers paper.
| predictions = model(x) | ||
|
|
||
| assert len(predictions) == 2 | ||
| assert sorted(predictions[0].keys()) == ["boxes", "labels", "masks", "scores"] |
There was a problem hiding this comment.
Should also check the shape of the masks.
| assert targets[0]["labels"].tolist() == [0] | ||
|
|
||
|
|
||
| def test_check_model(): |
There was a problem hiding this comment.
Redundant with test_predict_outputs_masks?
| area = torch.as_tensor(targets["area"]) | ||
| image_id = targets["image_id"] | ||
|
|
||
| # Filter out-of-bounds boxes and sync with transformed |
There was a problem hiding this comment.
Keep this comment?
| synced_metrics[key] = torch.tensor(value, device=self.device) | ||
| else: | ||
| synced_metrics[key] = value | ||
| elif self.model.task == "polygon": |
There was a problem hiding this comment.
Not obvious why this should be separate logic to boxes?
| results.update(self.trainer.logged_metrics) | ||
| results["predictions"] = self.predictions | ||
| if self.model.task == "box" or self.model.task == "point": | ||
| if self.model.task in ("box", "point", "polygon"): |
There was a problem hiding this comment.
Drop the condition here as this is valid for all model tasks
| return predictions.iloc[np.flatnonzero(kept)].reset_index(drop=True) | ||
|
|
||
|
|
||
| def reduce_polygons(predictions: pd.DataFrame, iou_threshold: float) -> pd.DataFrame: |
There was a problem hiding this comment.
Does this need to be a different function? Polygons aren't used at all here, only box bounds. If the result is that the dataframe/preds are filtered by index then it shouldn't matter if we ran reduce_boxes.
Might need to refactor with an output_columns argument?
|
Sorry, I actually didn't mean to trigger this into main yet. Thanks for looking at it. I thought I pushed to my fork. Let's hold this open and I will clean it up on monday. |
2 participants
Implement instance segmentation paralleling the box and point workflows:
Description
Related Issue(s)
AI-Assisted Development
AI tools used (if applicable):