JuliaWolleb / Diffusion-based-Segmentation

This is the official Pytorch implementation of the paper "Diffusion Models for Implicit Image Segmentation Ensembles".
MIT License
271 stars 35 forks source link

Input chnnels #51

Closed yug125lk closed 8 months ago

yug125lk commented 8 months ago

hi, thank you for sharing this great work,

I tried to experiment with the source code on my own dataset (one input image and one target image; images are in gray scale). I changed the script_util.py in the UNet class to: out_channels = 1 to have a grayscale output. return UNetModel( image_size=image_size, in_channels=2, model_channels=num_channels, out_channels=1,#(3 if not learn_sigma else 6),

    num_res_blocks=num_res_blocks,
    attention_resolutions=tuple(attention_ds),
    dropout=dropout,
    channel_mult=channel_mult,
    num_classes=(NUM_CLASSES if class_cond else None),
    use_checkpoint=use_checkpoint,
    use_fp16=use_fp16,
    num_heads=num_heads,
    num_head_channels=num_head_channels,
    num_heads_upsample=num_heads_upsample,
    use_scale_shift_norm=use_scale_shift_norm,
    resblock_updown=resblock_updown,
    use_new_attention_order=use_new_attention_order,
)

Then I received the following error: Traceback (most recent call last): File "/content/drive/MyDrive/seg/seg/scripts/segmentation_train.py", line 147, in main() File "/content/drive/MyDrive/seg/seg/scripts/segmentation_train.py", line 118, in main ).run_loop() File "/content/drive/MyDrive/seg/seg/./guided_diffusion/train_util.py", line 195, in run_loop self.run_step(batch, cond) File "/content/drive/MyDrive/seg/seg/./guided_diffusion/train_util.py", line 220, in run_step lossseg, losscls, lossrec, sample = self.forward_backward(batch, cond) File "/content/drive/MyDrive/seg/seg/./guided_diffusion/train_util.py", line 252, in forward_backward losses1 = compute_losses() File "/content/drive/MyDrive/seg/seg/./guided_diffusion/gaussian_diffusion.py", line 951, in training_losses_segmentation assert model_output.shape == (B, C 2, x_t.shape[2:]) AssertionError

YI-HAO-SU commented 8 months ago

You have to check more information about channels in "/content/drive/MyDrive/seg/seg/./guided_diffusion/gaussian_diffusion.py". I also adjusted the channel parameter to moderate the three channels' images, a lot of info about channels needs to be modified instead of only changing the output channel in Unet.py.

yug125lk commented 8 months ago

You have to check more information about channels in "/content/drive/MyDrive/seg/seg/./guided_diffusion/gaussian_diffusion.py". I also adjusted the channel parameter to moderate the three channels' images, a lot of info about channels needs to be modified instead of only changing the output channel in Unet.py.

Please, could you help me? What did you modify? I changed this line in gaussian_diffusion.py: img.shape != (1, 3, 128, 128): and it doesn't work. Could you provide me with the modified version of gaussian_diffusion.py?

JuliaWolleb commented 8 months ago

Hi Did you set the flag learn_sigma to True or False? If you set it to True, your model will predict the mean and the variance, so you will need two output channels.

yug125lk commented 8 months ago

It works. Thank you so much.

yug125lk commented 8 months ago

It works. Thank you so much.