Help with semantic segmentation

[lcoandrade]

I'm trying to perform an image semantic segmentation (segment mining fields) using lightning-flash on Colab. My images are all RGB/uint8/512x512 and the masks are L/uint8/512x512.

I have folders with train/train_masks, validation/validation_masks and test/test_masks data.

I would like to perform some transformations, like ToTensor, RandomHorizontalFlip and RandomVerticalFlip.

My code, so far is this:

!pip install lightning-flash
!pip install 'lightning-flash[image]'
!pip install torchmetrics

import torch
import flash
from flash.image import SemanticSegmentation, SemanticSegmentationData
import os
import torchmetrics
import ssl

from google.colab import drive

drive.mount("/content/drive")

DATA_DIR = '/content/drive/MyDrive/data/'

x_train_dir = os.path.join(DATA_DIR, 'train_images')
y_train_dir = os.path.join(DATA_DIR, 'train_masks')

x_valid_dir = os.path.join(DATA_DIR, 'val_images')
y_valid_dir = os.path.join(DATA_DIR, 'val_masks')

x_test_dir = os.path.join(DATA_DIR, 'test_images')
y_test_dir = os.path.join(DATA_DIR, 'test_masks')

datamodule = SemanticSegmentationData.from_folders(
    train_folder=x_train_dir,
    train_target_folder=y_train_dir,
    val_folder=x_valid_dir,
    val_target_folder=y_valid_dir,
    test_folder=x_test_dir,
    test_target_folder=y_test_dir,
    transform_kwargs=dict(image_size=(256, 256)),
    num_classes=1,
    batch_size=16,
)

#avoid certification error
ssl._create_default_https_context = ssl._create_unverified_context

model = SemanticSegmentation(
    head="unetplusplus", 
    backbone="efficientnet-b3", 
    pretrained="imagenet", 
    num_classes=datamodule.num_classes,
    metrics=[torchmetrics.Accuracy(), torchmetrics.Dice(), torchmetrics.F1Score(), torchmetrics.Recall()],
    optimizer='Adam',
    learning_rate=0.001
    )

model.configure_optimizers()

GPUS = torch.cuda.device_count()
EPOCHS = 2
if GPUS > 0:
    trainer = flash.Trainer(max_epochs=EPOCHS, gpus=GPUS)
else:
    trainer = flash.Trainer(max_epochs=EPOCHS)

trainer.finetune(model, datamodule=datamodule)
trainer.save_checkpoint(os.path.join(DATA_DIR, 'flash_trained.ckpt'))

But, when I run the code, I get this error:
IndexError Traceback (most recent call last)

<ipython-input-7-11e2ce087ca0> in <module>
      6 
      7 #trainer.fit(model, datamodule=datamodule)
----> 8 trainer.finetune(model, datamodule=datamodule, strategy="freeze")
      9 trainer.save_checkpoint("semantic_segmentation_model.pt")

19 frames
/usr/local/lib/python3.7/dist-packages/flash/core/trainer.py in finetune(self, model, train_dataloader, val_dataloaders, datamodule, strategy, train_bn)
    162         """
    163         self._resolve_callbacks(model, strategy, train_bn=train_bn)
--> 164         return super().fit(model, train_dataloader, val_dataloaders, datamodule)
    165 
    166     def predict(

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    695         self.strategy.model = model
    696         self._call_and_handle_interrupt(
--> 697             self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    698         )
    699 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _call_and_handle_interrupt(self, trainer_fn, *args, **kwargs)
    648                 return self.strategy.launcher.launch(trainer_fn, *args, trainer=self, **kwargs)
    649             else:
--> 650                 return trainer_fn(*args, **kwargs)
    651         # TODO(awaelchli): Unify both exceptions below, where `KeyboardError` doesn't re-raise
    652         except KeyboardInterrupt as exception:

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    735             ckpt_path, model_provided=True, model_connected=self.lightning_module is not None
    736         )
--> 737         results = self._run(model, ckpt_path=self.ckpt_path)
    738 
    739         assert self.state.stopped

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _run(self, model, ckpt_path)
   1166         self._checkpoint_connector.resume_end()
   1167 
-> 1168         results = self._run_stage()
   1169 
   1170         log.detail(f"{self.__class__.__name__}: trainer tearing down")

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _run_stage(self)
   1252         if self.predicting:
   1253             return self._run_predict()
-> 1254         return self._run_train()
   1255 
   1256     def _pre_training_routine(self):

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _run_train(self)
   1274 
   1275         with isolate_rng():
-> 1276             self._run_sanity_check()
   1277 
   1278         # enable train mode

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _run_sanity_check(self)
   1343             # run eval step
   1344             with torch.no_grad():
-> 1345                 val_loop.run()
   1346 
   1347             self._call_callback_hooks("on_sanity_check_end")

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/loop.py in run(self, *args, **kwargs)
    198             try:
    199                 self.on_advance_start(*args, **kwargs)
--> 200                 self.advance(*args, **kwargs)
    201                 self.on_advance_end()
    202                 self._restarting = False

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py in advance(self, *args, **kwargs)
    153         if self.num_dataloaders > 1:
    154             kwargs["dataloader_idx"] = dataloader_idx
--> 155         dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
    156 
    157         # store batch level output per dataloader

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/loop.py in run(self, *args, **kwargs)
    198             try:
    199                 self.on_advance_start(*args, **kwargs)
--> 200                 self.advance(*args, **kwargs)
    201                 self.on_advance_end()
    202                 self._restarting = False

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py in advance(self, data_fetcher, dl_max_batches, kwargs)
    141 
    142         # lightning module methods
--> 143         output = self._evaluation_step(**kwargs)
    144         output = self._evaluation_step_end(output)
    145 

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py in _evaluation_step(self, **kwargs)
    238         """
    239         hook_name = "test_step" if self.trainer.testing else "validation_step"
--> 240         output = self.trainer._call_strategy_hook(hook_name, *kwargs.values())
    241 
    242         return output

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/trainer/trainer.py in _call_strategy_hook(self, hook_name, *args, **kwargs)
   1704 
   1705         with self.profiler.profile(f"[Strategy]{self.strategy.__class__.__name__}.{hook_name}"):
-> 1706             output = fn(*args, **kwargs)
   1707 
   1708         # restore current_fx when nested context

/usr/local/lib/python3.7/dist-packages/pytorch_lightning/strategies/strategy.py in validation_step(self, *args, **kwargs)
    368         with self.precision_plugin.val_step_context():
    369             assert isinstance(self.model, ValidationStep)
--> 370             return self.model.validation_step(*args, **kwargs)
    371 
    372     def test_step(self, *args: Any, **kwargs: Any) -> Optional[STEP_OUTPUT]:

/usr/local/lib/python3.7/dist-packages/flash/image/segmentation/model.py in validation_step(self, batch, batch_idx)
    151     def validation_step(self, batch: Any, batch_idx: int) -> Any:
    152         batch = (batch[DataKeys.INPUT], batch[DataKeys.TARGET])
--> 153         return super().validation_step(batch, batch_idx)
    154 
    155     def test_step(self, batch: Any, batch_idx: int) -> Any:

/usr/local/lib/python3.7/dist-packages/flash/core/model.py in validation_step(self, batch, batch_idx)
    423 
    424     def validation_step(self, batch: Any, batch_idx: int) -> None:
--> 425         output = self.step(batch, batch_idx, self.val_metrics)
    426         log_kwargs = {"batch_size": output.get(OutputKeys.BATCH_SIZE, None)} if _PL_GREATER_EQUAL_1_5_0 else {}
    427         self.log_dict(

/usr/local/lib/python3.7/dist-packages/flash/core/model.py in step(self, batch, batch_idx, metrics)
    360         output = {OutputKeys.OUTPUT: y_hat}
    361         y_hat = self.to_loss_format(output[OutputKeys.OUTPUT])
--> 362         losses = {name: l_fn(y_hat, y) for name, l_fn in self.loss_fn.items()}
    363 
    364         y_hat = self.to_metrics_format(output[OutputKeys.OUTPUT])

/usr/local/lib/python3.7/dist-packages/flash/core/model.py in <dictcomp>(.0)
    360         output = {OutputKeys.OUTPUT: y_hat}
    361         y_hat = self.to_loss_format(output[OutputKeys.OUTPUT])
--> 362         losses = {name: l_fn(y_hat, y) for name, l_fn in self.loss_fn.items()}
    363 
    364         y_hat = self.to_metrics_format(output[OutputKeys.OUTPUT])

/usr/local/lib/python3.7/dist-packages/torch/nn/functional.py in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3012     if size_average is not None or reduce is not None:
   3013         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3014     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
   3015 
   3016 

IndexError: Target 255 is out of bounds.

How can I fix that?

[aniketmaurya]

@krshrimali