Refactor image name handling in SegmentationDOFA test step

geo_deep_learning/tasks_with_models/segmentation_dofa.py

@@ -87,7 +87,6 @@ def validation_step(self, batch, batch_idx):
 
     def test_step(self, batch, batch_idx):
         x = batch["image"]
-        image_names = batch["image_name"]
         y = batch["mask"]
         y = y.squeeze(1).long()
         y_hat = self(x)
@@ -101,7 +100,7 @@ def test_step(self, batch, batch_idx):
             num_samples = min(remaining_samples, len(x))
             for i in range(num_samples):
                 image = x[i]
-                image_name = image_names[i]
+                image_name = batch["image_name"][i]
                 image = denormalization(image, mean=self.mean, std=self.std, data_type_max=self.data_type_max)
                 fig = visualize_prediction(image,
                                             y[i],