[segmentation.models] Update models to handle updated samples

src/mednet/libs/segmentation/models/driu.py

@@ -161,23 +161,23 @@ class DRIU(Model):
             super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(self._augmentation_transforms(images))
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(images)
         return self._validation_loss(outputs, ground_truths, masks)
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        output = self(batch[0])[1]
+        output = self(batch[0]["image"])[1]
         return torch.sigmoid(output)
 
     def configure_optimizers(self):

src/mednet/libs/segmentation/models/driu_bn.py

@@ -164,23 +164,23 @@ class DRIUBN(Model):
             super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(self._augmentation_transforms(images))
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(images)
         return self._validation_loss(outputs, ground_truths, masks)
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        output = self(batch[0])[1]
+        output = self(batch[0]["image"])[1]
         return torch.sigmoid(output)
 
     def configure_optimizers(self):

src/mednet/libs/segmentation/models/driu_od.py

@@ -146,23 +146,23 @@ class DRIUOD(Model):
             super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(self._augmentation_transforms(images))
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(images)
         return self._validation_loss(outputs, ground_truths, masks)
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        output = self(batch[0])[1]
+        output = self(batch[0]["image"])[1]
         return torch.sigmoid(output)
 
     def configure_optimizers(self):

src/mednet/libs/segmentation/models/driu_pix.py

@@ -150,23 +150,23 @@ class DRIUPix(Model):
             super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(self._augmentation_transforms(images))
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(images)
         return self._validation_loss(outputs, ground_truths, masks)
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        output = self(batch[0])[1]
+        output = self(batch[0]["image"])[1]
         return torch.sigmoid(output)
 
     def configure_optimizers(self):

src/mednet/libs/segmentation/models/hed.py

@@ -165,23 +165,23 @@ class HED(Model):
             super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(self._augmentation_transforms(images))
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(images)
         return self._validation_loss(outputs, ground_truths, masks)
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        output = self(batch[0])[1]
+        output = self(batch[0]["image"])[1]
         return torch.sigmoid(output)
 
     def configure_optimizers(self):

src/mednet/libs/segmentation/models/lwnet.py

@@ -366,23 +366,23 @@ class LittleWNet(Model):
         return x1, x2
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(self._augmentation_transforms(images))
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(images)
         return self._validation_loss(outputs, ground_truths, masks)
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        output = self(batch[0])[1]
+        output = self(batch[0]["image"])[1]
         return torch.sigmoid(output)
 
     def configure_optimizers(self):

src/mednet/libs/segmentation/models/m2unet.py

@@ -213,23 +213,23 @@ class M2UNET(Model):
             super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(self._augmentation_transforms(images))
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(images)
         return self._validation_loss(outputs, ground_truths, masks)
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        output = self(batch[0])[1]
+        output = self(batch[0]["image"])[1]
         return torch.sigmoid(output)
 
     def configure_optimizers(self):

src/mednet/libs/segmentation/models/unet.py

@@ -154,23 +154,23 @@ class Unet(Model):
             super().set_normalizer(dataloader)
 
     def training_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(self._augmentation_transforms(images))
         return self._train_loss(outputs, ground_truths, masks)
 
     def validation_step(self, batch, batch_idx):
-        images = batch[0]
-        ground_truths = batch[1]["target"]
-        masks = batch[1]["mask"]
+        images = batch[0]["image"]
+        ground_truths = batch[0]["target"]
+        masks = batch[0]["mask"]
 
         outputs = self(images)
         return self._validation_loss(outputs, ground_truths, masks)
 
     def predict_step(self, batch, batch_idx, dataloader_idx=0):
-        output = self(batch[0])[1]
+        output = self(batch[0]["image"])[1]
         return torch.sigmoid(output)
 
     def configure_optimizers(self):