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
79 stars 4 forks source link

Some of the generated images are normal and some contain noise spots #12

Open fxsly opened 1 month ago

fxsly commented 1 month ago

Hello, when I trained the model with my own dataset(MRI slices) about 1000 npy files, I encountered the problem above. According to your updated code, I transformed my npy files to PIL images , the detailed code is as follows:

`def npy_to_pil(file):

np_array=np.load(file)

np_img=normalize_data(np_array)#0-255

np_img=np_img.round().astype('uint8')

pil_images=Image.fromarray(np_img.squeeze(0),mode='L')

return pil_images

def normalize_data(data):

min=data.min()

max=data.max()

data=(data-min)/(max-min)*255.0

return  data`

And my settings are as follows: CUDA_VISIBLE_DEVICES={DEVICES} python3 main.py --mode train --model_type DDPM --img_size 256 --num_img_channels 1 --dataset {DATASET_NAME} --train_batch_size 8 --eval_batch_size 8 --eval_sample_size 1000 --segmentation_guided Here are my results and mask used: 图片1 图片2

I have some guesses: First is the number of my dataset. Is it possible that the number of my dataset is too small? Second is the type of my dataset. My dataset is MRI, is MRI training difficultly? Because MRI is little more comlex compared to CT. Should I train my dataset with more epochs? Third is the range of the data. When I trained my dataset with the code at first for 400 epochs, the range of my data is [-1,1 ], the trained result is as follows: 微信图片_20240730172008 It looks worse than the PIL image. So does the range of data have bad influence on the result at the same epochs? If I use [-1,1] data range, I trained the model with more epochs, will the influence vanish and generate images with no noises? Hope to get your response.

nickk124 commented 1 month ago

Hi!

Generating noisy images happens sometimes if you don't train for long enough, in my experience, and also other users' (e.g. https://github.com/mazurowski-lab/segmentation-guided-diffusion/issues/11). While 400 epochs is a lot, your dataset is somewhat small, so the total number of images seen in training may need to be higher (for example, when I trained on ~10,000 images, I saw noisy images generated as late as epoch 200). So, my recommendation is just to train for longer: try maybe 600 or 800 epochs total, to start?

MRI may be a bit harder to learn than CT, but I think training for longer should improve things. In terms of the data range, the code automatically normalizes the images when they are converted into torch tensors, so this shouldn't be an issue.

Does that help?

fxsly commented 1 month ago

Hi!

Generating noisy images happens sometimes if you don't train for long enough, in my experience, and also other users' (e.g. #11). While 400 epochs is a lot, your dataset is somewhat small, so the total number of images seen in training may need to be higher (for example, when I trained on ~10,000 images, I saw noisy images generated as late as epoch 200). So, my recommendation is just to train for longer: try maybe 600 or 800 epochs total, to start?

MRI may be a bit harder to learn than CT, but I think training for longer should improve things. In terms of the data range, the code automatically normalizes the images when they are converted into torch tensors, so this shouldn't be an issue.

Does that help?

Thank you for your reply! I tried to train the model for 800 epochs,but there was still noisy images. Maybe I should adjust the parameters?

nickk124 commented 1 month ago

Did training for longer decrease the typical noise, even if it's still present overall? If so, because your dataset is small, then it may work to train for even longer; you might as well try another 800 epochs. Otherwise, yes, hyperparameter adjustment may help!

fxsly commented 1 month ago

Did training for longer decrease the typical noise, even if it's still present overall? If so, because your dataset is small, then it may work to train for even longer; you might as well try another 800 epochs. Otherwise, yes, hyperparameter adjustment may help!

Thank you! I think the result of 800 epochs is the same as 400 epochs'. 1723517410795 Maybe I should adjust the hyperparameter or train more epochs! Thank you for your advice!

nickk124 commented 1 month ago

Hi, yeah I think hyperparameter adjustment may be the right idea -- I would suggest reducing the learning rate, perhaps to 2e-5? You can change it at https://github.com/mazurowski-lab/segmentation-guided-diffusion/blob/b1ef8b137eaaefab0210e52b4c49f34ff6067fa6/training.py#L29

Please let me know if that works!