Pytorch-cyclegan-and-pix2pix: more than one input image

[armaneshaghi]

Perhaps a naive question, but I am wondering whether I can use this repository for multiple input images and one output image (say A1, A2, A3 ==> B ) . My use case is in medical imaging where I get multiple modalities (e.g, A1 and A2) as input to a GAN to synthesise a new output (B).

[junyanz]

Yes. You can write a custom data loader. In this data loader, you need to load three images from the disk separately, concatenate them, and feed them as input to the generator. You need to also change --input_nc. See here on how to implement your own data loader. Also, see the overview of this repo.

[villanuevab]

Would you suggest stacking them along the C dimension, then? Since we're modifying input_nc. That is, for A1, A2, A3, where A1 has dim (H,W,C), A2 has dim (H,W,C), and A3 has dim (H,W,C), we stack these along the C dimension such that the input A123 has dim (H,W,3C). And the output B will be (H,W,C). Is this understanding correct?

[junyanz]

Yes. set input_nc=3C.

[villanuevab]

Update: resolved this. The error was in the dict returned by __getitem__. Instead of:

return {'data_A': data_A, 'data_B': data_B, 'A_paths': A_paths, 'B_paths': B_path}

I now have:

return {'A': data_A, 'B': data_B, 'A_paths': A_paths, 'B_paths': B_path}

---

Currently I'm able to load images properly using a custom data loader __init__. My __getitem__ seems to be working as well. However, I'm getting errors training the pix2pix model. I receive the following error:

Traceback (most recent call last):
  File "train.py", line 50, in <module>
    model.set_input(data)         # unpack data from dataset and apply preprocessing
  File "/home/blancavillanueva/pytorch-CycleGAN-and-pix2pix/models/pix2pix_model.py", line 82, in set_input
    self.real_A = input['A' if AtoB else 'B'].to(self.device)
KeyError: 'A'

I have multiple paths for the A image, but only a single path for the B image. I assign them in __getitem__ like this:

return {'data_A': data_A, 'data_B': data_B, 'A_paths': A_paths, 'B_paths': B_path}

Is it possible to keep my current directory structure (which is different from the aligned dataset's expected /path/to/data/A/train and /path/to/data/B/train) or do I need to move my files into directories to match the expected directory structure? The files are quite large and moving them around will be expensive.

[LPUoO]

Hello @junyanz, @taesungp, @SsnL,

Thank you very much for making this amazing tool and making it readily available.

I am also interested in using 2 or 3 images to predict a single output.
When using 3 greyscale images would you concatenate them or would you set each image as R, G or B?
When using 2 greyscale images would you concatenate them or would you set each image as R, G and set B to zero?

I guess that setting the images as RGB will reduce computational cost but I am ensure if that's okay for pix2pix?

Thank you very much for your help