Getting the error while converting the tensors to Images using cv2

1

Hi,

I have converted the tensors which I’m getting as an output(test set) from my UNet to images. When I’m testing the model it is throwing out of index exception which I’m not getting when I didn’t convert the tensors to images. I have 74 images in my test set but after getting tested with one image it is throwing out of index exception in my custom dataset class. any idea will there will be any problem if I do a conversion like that?

# this is where im converting
 def test_step(self, batch, batch_nb):
        x, y = batch
        y_hat = self.forward(x)
        # print(y_hat.size())
        img = y_hat.cpu().detach().numpy()
        directory = r'D:\Mobile-research project\dataset\gaussian_blur\target_outputs'
        os.chdir(directory)
        print(os.listdir(directory))
        filename = 'savedImage.jpg'
        cv2.imwrite(filename, img)
        loss = torch.nn.MSELoss()
        op_loss = loss(y_hat, y)
        return {'test_loss': op_loss}

#custom dataset

    def __getitem__(self, i):
        idx = self.ids[i]

        img_files = glob.glob(os.path.join(self.img_dir, idx+'.*'))
        mask_files = glob.glob(os.path.join(self.mask_dir, idx+'.*'))

        assert len(img_files) == 1, f'{idx}: {img_files}'
        assert len(mask_files) == 1, f'{idx}: {mask_files}'

        # use Pillow's Image to read .gif mask
        # https://answers.opencv.org/question/185929/how-to-read-gif-in-python/
        img = Image.open(img_files[0])
        mask = Image.open(mask_files[0])
        # to_tensor = transforms.ToTensor()
        # img_t = to_tensor(img)
        # a = img_t.shape
        assert img.size == mask.size, f'{img.shape} # {mask.shape}'

        return self.transforms(img),\
               self.transforms(mask)

I’m getting this error:

File “D:\unet-lightning-master -original\unet-lightning-master-original\dataset.py”, line 35, in getitem
assert len(img_files) == 1, f’{idx}: {img_files}’
AssertionError: 0101434: []
Testing: 1%|▏ | 1/72 [00:00<00:40, 1.77it/s]

Process finished with exit code 1