manually add_image to subject will cause dataloader stuck

[jxchen01]

Is there an existing issue for this?

• [X] I have searched the existing issues

Problem summary

manually add_image (created by numpy array, not file path) causes dataloader stuck forever.

Code for reproduction

# code 1
self.subjects = []
for ds in dataset_list:
    subject = tio.Subject(
        source=tio.ScalarImage(ds["filename1"]),
        target=tio.LabelMap(ds["filename2"]),
        center_image=tio.LabelMap(ds["filename3"]),
        class_image=tio.LabelMap(ds["filename4"]),
    )
    self.subjects.append(subject)
my_dataloader1 = tio.SubjectsDataset(self.subjects)

# code 2
self.subjects = []
for ds in dataset_list:
    subject = tio.Subject(
        source=tio.ScalarImage(ds["filename1"]),
        target=tio.LabelMap(ds["filename2"]),
    )
    im = subject["target"][tio.DATA]

    # do some processing with func_foo1 and func_foo2
    center_image = func_foo1(im)
    class_image = func_foo2(im)

    # manually add these two generated images to subject
    subject.add_image(tio.LabelMap(tensor=center_image), "center_image")
    subject.add_image(tio.LabelMap(tensor=class_image), "class_image")

    self.subjects.append(subject)
my_dataloader2 = tio.SubjectsDataset(self.subjects)

Actual outcome

code 1 above can run without any problem. But, if I want to generate my "center_image" and "class_image" on-the-fly, instead of pre-saved them in files, and then manually add those generated "center_image" and "class_image" to the subject, the created "my_dataloader2" will stuck. By "stuck", I mean the code cannot enter the forward step, and stuck somewhere in the dataloader step (hard to determine where it stucks).

Error messages

no error message, code 2 just stucks.

Expected outcome

I expect code 1 and code 2 should work similarly.

System info

Platform:   Linux-5.4.0-89-generic-x86_64-with-glibc2.17
TorchIO:    0.18.76
PyTorch:    1.8.2
SimpleITK:  2.1.1 (ITK 5.2)
NumPy:      1.21.2
Python:     3.8.12 (default, Oct 12 2021, 13:49:34) 
[GCC 7.5.0]

[fepegar]

Hi, @jxchen01. Thanks for reporting. Can you please try to provide a reproducible example?

[jxchen01]

Currently, it is living in a large code base. I will try to make a short reproducible example

[LeonieBorne]

Hi, I have a problem that seems similar: with the code below, my dataloader is stuck (but only when I use the plot() function before). I don't have this problem if I don't transform the images. I hope that helps, let me know if you need more information, thanks!

Code for reproduction

import torchio as tio
import time
import torch
from pathlib import Path


dataset_dir = Path('/path/to/data/ixi_tiny')
images_dir = dataset_dir / 'image'
labels_dir = dataset_dir / 'label'
image_paths = sorted(images_dir.glob('*.nii.gz'))
label_paths = sorted(labels_dir.glob('*.nii.gz'))
assert len(image_paths) == len(label_paths)

subjects = []
for (image_path, label_path) in zip(image_paths, label_paths):
    subject = tio.Subject(
        mri=tio.ScalarImage(image_path),
        brain=tio.LabelMap(label_path),
    )
    subjects.append(subject)

dataset = tio.SubjectsDataset(
        subjects, transform=tio.RandomAffine())  
        
loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=1,
    shuffle=True,
    num_workers=1,
)

start_time = time.time()
one_batch = next(iter(loader))
print(f'Load one batch in {time.time()-start_time:.2f} s.')

start_time = time.time()
dataset[0].plot()
print(f'Plot subject in {time.time()-start_time:.2f} s.')

start_time = time.time()
one_batch = next(iter(loader))
print(f'Load one batch in {time.time()-start_time:.2f} s.')

Actual outcome

Load one batch in 0.42 s.
Plot subject in 1.14 s.

System info

Platform:   Linux-3.10.0-693.17.1.el7.x86_64-x86_64-with-glibc2.17
TorchIO:    0.18.82
PyTorch:    1.11.0
SimpleITK:  2.0.2 (ITK 5.1)
NumPy:      1.21.5
Python:     3.9.12 (main, Apr  5 2022, 06:56:58) 
[GCC 7.5.0]

[fepegar]

Closing stale issue. Feel free to reopen if you can provide a MWE.