Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Trainers: add Instance Segmentation Task #2513

Merged
merged 58 commits into from
Feb 25, 2025
Merged
Changes from 1 commit
Commits
Show all changes
58 commits
Select commit Hold shift + click to select a range
d00c087
Add files via upload
ariannasole23 Jan 9, 2025
52daa1c
Add files via upload
ariannasole23 Jan 9, 2025
68756a7
Update instancesegmentation.py
ariannasole23 Jan 9, 2025
e249883
Merge branch 'microsoft:main' into main
ariannasole23 Jan 20, 2025
7676ac3
Update and rename instancesegmentation.py to instance_segmentation.py
ariannasole23 Jan 20, 2025
0fa7b07
Update test_instancesegmentation.py
ariannasole23 Jan 21, 2025
b4334f0
Update instance_segmentation.py
ariannasole23 Jan 21, 2025
a160baa
Update __init__.py
ariannasole23 Jan 21, 2025
fa8697b
Update instance_segmentation.py
ariannasole23 Jan 21, 2025
f6ceed1
Update instance_segmentation.py
ariannasole23 Jan 27, 2025
619760b
Add files via upload
ariannasole23 Jan 27, 2025
d9158a0
Update test_instancesegmentation.py
ariannasole23 Jan 27, 2025
9f48f50
Update and rename test_instancesegmentation.py to test_trainer_instan…
ariannasole23 Jan 28, 2025
63aefc8
Update instance_segmentation.py
ariannasole23 Jan 28, 2025
70074e7
Add files via upload
ariannasole23 Jan 28, 2025
b3de001
Creato con Colab
ariannasole23 Jan 28, 2025
d70f1e3
Creato con Colab
ariannasole23 Jan 28, 2025
1e68d2d
Creato con Colab
ariannasole23 Jan 28, 2025
98c836a
Merge branch 'microsoft:main' into main
ariannasole23 Feb 5, 2025
9664834
Update instance_segmentation.py
ariannasole23 Feb 5, 2025
f802574
Delete test_trainer.ipynb
ariannasole23 Feb 5, 2025
3c86306
Delete test_trainer_instancesegmentation.py
ariannasole23 Feb 5, 2025
7ec3930
Update and rename test_instancesegmentation.py to test_instance_segme…
ariannasole23 Feb 5, 2025
927f7fc
Update instance_segmentation.py
ariannasole23 Feb 5, 2025
4f1cecf
Update test_instance_segmentation.py
ariannasole23 Feb 5, 2025
21e0af2
Update instance_segmentation.py
ariannasole23 Feb 6, 2025
3956d23
Update instance_segmentation.py
ariannasole23 Feb 6, 2025
0e458a5
Update instance_segmentation.py run ruff
ariannasole23 Feb 6, 2025
870845b
Merge remote-tracking branch 'upstream/main'
adamjstewart Feb 20, 2025
fafb001
Ruff
adamjstewart Feb 20, 2025
ad7197d
dos2unix
adamjstewart Feb 20, 2025
954e898
Add support for MSI, weights
adamjstewart Feb 20, 2025
3c6ee68
Update tests
adamjstewart Feb 20, 2025
7c4e30c
timm and torchvision are not compatible
adamjstewart Feb 20, 2025
7c34d4a
Finalize trainer code, simpler
adamjstewart Feb 21, 2025
649a877
Update VHR10 tests
adamjstewart Feb 21, 2025
4f201fd
Uniformity
adamjstewart Feb 21, 2025
006cfa9
Fix most tests
adamjstewart Feb 21, 2025
b3a4e44
100% coverage
adamjstewart Feb 21, 2025
1d80adc
Fix datasets tests
adamjstewart Feb 21, 2025
d8e8fe6
Fix weight tests
adamjstewart Feb 21, 2025
f774875
Fix MSI support
adamjstewart Feb 21, 2025
c823fd0
Fix parameter replacement
adamjstewart Feb 21, 2025
94e8001
Fix minimum tests
adamjstewart Feb 21, 2025
5e01c96
Fix minimum tests
adamjstewart Feb 21, 2025
2460b26
Add all unpacked data
adamjstewart Feb 21, 2025
d63cf85
Fix tests
adamjstewart Feb 21, 2025
f85a72e
Undo FTW changes
adamjstewart Feb 21, 2025
683c162
Undo FTW changes
adamjstewart Feb 21, 2025
8a9c0e9
Undo FTW changes
adamjstewart Feb 21, 2025
b072a38
Remove dead code
adamjstewart Feb 21, 2025
c4b5d17
Remove dead code, match detection style
adamjstewart Feb 22, 2025
801c0ba
Try newer torchmetrics
adamjstewart Feb 22, 2025
4640d6c
Try newer torchmetrics
adamjstewart Feb 22, 2025
1d2a595
Try newer torchmetrics
adamjstewart Feb 22, 2025
8f165ab
More metrics
adamjstewart Feb 23, 2025
7b6182d
Fix mypy
adamjstewart Feb 23, 2025
335f072
Fix and test weights=True, num_classes!=91
adamjstewart Feb 25, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Update instance_segmentation.py
ariannasole23 authored Jan 27, 2025

Verified

This commit was signed with the committer’s verified signature.
woodruffw William Woodruff
commit f6ceed184ecc50a2af487ad3c2586e541dc8abfe
138 changes: 112 additions & 26 deletions torchgeo/trainers/instance_segmentation.py
Original file line number Diff line number Diff line change
@@ -3,20 +3,20 @@

"""Trainers for instance segmentation."""

from typing import Any
import torch.nn as nn
import torch
from torch import Tensor
from torchmetrics.detection.mean_ap import MeanAveragePrecision
from torchmetrics import MetricCollection
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection import maskrcnn_resnet50_fpn, MaskRCNN_ResNet50_FPN_Weights
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from .base import BaseTask

import matplotlib.pyplot as plt
from matplotlib.figure import Figure
from ..datasets import RGBBandsMissingError, unbind_samples

Check failure on line 18 in torchgeo/trainers/instance_segmentation.py

GitHub Actions / ruff

Ruff (I001)

torchgeo/trainers/instance_segmentation.py:6:1: I001 Import block is un-sorted or un-formatted


class InstanceSegmentationTask(BaseTask):
"""Instance Segmentation."""

@@ -25,7 +25,7 @@
model: str = 'mask_rcnn',
backbone: str = 'resnet50',
weights: str | bool | None = None,
num_classes: int = 2,
num_classes: int = 2,
lr: float = 1e-3,
patience: int = 10,
freeze_backbone: bool = False,
@@ -49,12 +49,12 @@
"""
self.weights = weights
super().__init__()
# self.save_hyperparameters()
# self.model = None
# self.validation_outputs = []
# self.test_outputs = []
# self.configure_models()
# self.configure_metrics()
self.save_hyperparameters()
self.model = None
self.validation_outputs = []
self.test_outputs = []
self.configure_models()
self.configure_metrics()

def configure_models(self) -> None:
"""Initialize the model.
@@ -67,11 +67,12 @@

if model == 'mask_rcnn':
# Load the Mask R-CNN model with a ResNet50 backbone
self.model = maskrcnn_resnet50_fpn(weights=self.weights is True)
self.model = maskrcnn_resnet50_fpn(weights=MaskRCNN_ResNet50_FPN_Weights.DEFAULT)

# Update the classification head to predict `num_classes`
in_features = self.model.roi_heads.box_predictor.cls_score.in_features
self.model.roi_heads.box_predictor = nn.Linear(in_features, num_classes)
# self.model.roi_heads.box_predictor = nn.Linear(in_features, num_classes)
self.model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

# Update the mask head for instance segmentation
in_features_mask = self.model.roi_heads.mask_predictor.conv5_mask.in_channels
@@ -126,13 +127,50 @@
Updates metrics and stores predictions/targets for further analysis.
"""
images, targets = batch['image'], batch['target']
outputs = self.model(images)
self.metrics.update(outputs, targets)
self.validation_outputs.append((outputs, targets))

metrics_dict = self.metrics.compute()
self.log_dict(metrics_dict)
self.metrics.reset()
batch_size = images.shape[0]

outputs = self.model(images)

for target in targets:
target["masks"] = (target["masks"] > 0).to(torch.uint8)
target["boxes"] = target["boxes"].to(torch.float32)
target["labels"] = target["labels"].to(torch.int64)

# Compute the loss and predictions
loss_dict = self.model(images, targets) # list of dictionaries

print('\nDEBUG TRAINING LOSS\n')
print(f"Training loss: {loss_dict}")

# Post-process `loss_dict` to compute total loss
total_loss = 0.0
for loss in loss_dict:
if isinstance(loss, dict):
for key, value in loss.items():
# Ensure the loss component is a scalar tensor
if value.ndim == 0:
total_loss += value
else:
print(f"Skipping non-scalar loss: {key}, shape: {value.shape}")

# Post-process the outputs to ensure masks are in the correct format
for output in outputs:
if "masks" in output:
output["masks"] = (output["masks"] > 0.5).squeeze(1).to(torch.uint8)

# Sum the losses
self.log('val_loss', total_loss, batch_size=batch_size)

metrics = self.val_metrics(outputs, targets)
# Log only scalar values from metrics
scalar_metrics = {}
for key, value in metrics.items():
if isinstance(value, torch.Tensor) and value.numel() > 1:
# Cast to float if integer and compute mean
value = value.to(torch.float32).mean()
scalar_metrics[key] = value

self.log_dict(scalar_metrics, batch_size=batch_size)

# check
if (
@@ -162,21 +200,61 @@
f'image/{batch_idx}', fig, global_step=self.global_step
)
plt.close()


def test_step(self, batch: Any, batch_idx: int) -> None:
"""Compute the test loss and additional metrics."""

Check failure on line 205 in torchgeo/trainers/instance_segmentation.py

GitHub Actions / ruff

Ruff (D202)

torchgeo/trainers/instance_segmentation.py:205:9: D202 No blank lines allowed after function docstring (found 1)

images, targets = batch['image'], batch['target']
batch_size = images.shape[0]

outputs = self.model(images)
self.metrics.update(outputs, targets)
self.test_outputs.append((outputs, targets))

metrics_dict = self.metrics.compute()
self.log_dict(metrics_dict)
print('\nDEBUG THE PREDICTIONS\n')
print(f"Predictions for batch {batch_idx}: {outputs}")
print(f"Ground truth for batch {batch_idx}: {targets}")

for target in targets:
target["masks"] = target["masks"].to(torch.uint8)
target["boxes"] = target["boxes"].to(torch.float32)
target["labels"] = target["labels"].to(torch.int64)

for output in outputs:
output["masks"] = (output["masks"] > 0.5).squeeze(1).to(torch.uint8)

loss_dict = self.model(images, targets) # Compute all losses

# Post-process `loss_dict` to compute total loss
total_loss = 0.0
for loss in loss_dict:
if isinstance(loss, dict):
for key, value in loss.items():
# Ensure the loss component is a scalar tensor
if value.ndim == 0:
total_loss += value
else:
print(f"Skipping non-scalar loss: {key}, shape: {value.shape}")


# Sum the losses
self.log('test_loss', total_loss, batch_size=batch_size)

metrics = self.val_metrics(outputs, targets)
# Log only scalar values from metrics
scalar_metrics = {}
for key, value in metrics.items():
if isinstance(value, torch.Tensor) and value.numel() > 1:
# Cast to float if integer and compute mean
value = value.to(torch.float32).mean()
scalar_metrics[key] = value

def predict_step(self, batch: Any, batch_idx: int) -> Tensor:
self.log_dict(scalar_metrics, batch_size=batch_size)

print('\nDEBUG CAL METRICS\n')
print(f"Validation metrics: {metrics}")

return outputs

def predict_step(self, batch: Any, batch_idx: int) -> Any:
"""Perform inference on a batch of images.
Args:
@@ -185,7 +263,15 @@
Returns:
Predicted masks and bounding boxes for the batch.
"""
self.model.eval()
images = batch['image']
y_hat: Tensor = self.model(images)
return y_hat
outputs = self.model(images)
return outputs









Unchanged files with check annotations Beta

import torch
import pytorch_lightning as pl
from pytorch_lightning import LightningModule

Check failure on line 3 in tests/trainers/test_instancesegmentation.py

GitHub Actions / ruff

Ruff (F401)

tests/trainers/test_instancesegmentation.py:3:31: F401 `pytorch_lightning.LightningModule` imported but unused
from torch.utils.data import DataLoader
from torchgeo.datasets import VHR10
from torchgeo.main import main

Check failure on line 6 in tests/trainers/test_instancesegmentation.py

GitHub Actions / ruff

Ruff (F401)

tests/trainers/test_instancesegmentation.py:6:27: F401 `torchgeo.main.main` imported but unused
from torchgeo.trainers import InstanceSegmentationTask

Check failure on line 8 in tests/trainers/test_instancesegmentation.py

GitHub Actions / ruff

Ruff (I001)

tests/trainers/test_instancesegmentation.py:1:1: I001 Import block is un-sorted or un-formatted
# Custom collate function for DataLoader (required for Mask R-CNN models)
def collate_fn(batch):

Check failure on line 12 in tests/trainers/test_instancesegmentation.py

GitHub Actions / ruff

Ruff (ANN201)

tests/trainers/test_instancesegmentation.py:12:5: ANN201 Missing return type annotation for public function `collate_fn`

Check failure on line 12 in tests/trainers/test_instancesegmentation.py

GitHub Actions / ruff

Ruff (ANN001)

tests/trainers/test_instancesegmentation.py:12:16: ANN001 Missing type annotation for function argument `batch`
return tuple(zip(*batch))
# Initialize the VHR10 dataset
"""TorchGeo trainers."""
from .base import BaseTask
from .byol import BYOLTask
from .classification import ClassificationTask, MultiLabelClassificationTask
from .detection import ObjectDetectionTask
from .iobench import IOBenchTask

Check failure on line 10 in torchgeo/trainers/__init__.py

GitHub Actions / ruff

Ruff (F401)

torchgeo/trainers/__init__.py:10:22: F401 `.iobench.IOBenchTask` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
from .moco import MoCoTask
from .regression import PixelwiseRegressionTask, RegressionTask
from .segmentation import SemanticSegmentationTask
from .simclr import SimCLRTask
from .instance_segmentation import InstanceSegmentationTask

Check failure on line 15 in torchgeo/trainers/__init__.py

GitHub Actions / ruff

Ruff (I001)

torchgeo/trainers/__init__.py:6:1: I001 Import block is un-sorted or un-formatted

Check failure on line 15 in torchgeo/trainers/__init__.py

GitHub Actions / ruff

Ruff (F401)

torchgeo/trainers/__init__.py:15:36: F401 `.instance_segmentation.InstanceSegmentationTask` imported but unused; consider removing, adding to `__all__`, or using a redundant alias
__all__ = (
'BYOLTask',