search

mazurowski-lab / segmentation-guided-diffusion

[MICCAI 2024] Easy diffusion models (optionally with segmentation guidance) for medical images and beyond.
https://arxiv.org/abs/2402.05210
Other
56 stars 4 forks source link

No dataset preprocessing in segmentation guided mode #7

Closed lukalafaye closed 3 weeks ago

lukalafaye commented 3 weeks ago

You should add these two lines in main.py after line 237 because the dataset is not being processed using transform when using --segmentation_guided .

dataset_train.set_transform(transform) dataset_eval.set_transform(transform)

Otherwise when running main.py you will get this error: Traceback (most recent call last): File "main.py", line XXX, in main( File "main.py", line XXX, in main train_loop( File "/home/diffusion_models3/segmentation-guided-diffusion/training.py", line 87, in train_loop for step, batch in enumerate(train_dataloader): TypeError: default_collate: batch must contain tensors, numpy arrays, numbers, dicts or lists; found <class 'PIL.PngImagePlugin.PngImageFile'>

nickk124 commented 3 weeks ago

Thanks for the catch, I just implemented the change in the latest commit. Cheers!