search

junyanz / pytorch-CycleGAN-and-pix2pix

Image-to-Image Translation in PyTorch
Other
22.8k stars 6.29k forks source link

Training pix2pix with High resolution images #1516

Closed maryam95hallal closed 1 year ago

maryam95hallal commented 1 year ago

Greetings Thanks for your amazing work! I have one issue and one problem I fixed to report ,and a question to ask. I'm using google collab to train pix2pix model, I need high resolution output so I tested the recommended settings --preprocess scale_width_and_crop --load_size 1024 --crop_size 360 and it worked for training cyclegan but when training pix2pix network I got the following error: RuntimeError: Sizes of tensors must match except in dimension 1. Expected size 5 but got size 4 for tensor number 1 in the list. This is my full training command: !python train.py --dataroot /content/gdrive/MyDrive/pix2pix_official/dataset/dataset --name watermark_remover_pix2pix_1024 --model pix2pix --preprocess scale_width_and_crop --load_size 1024 --crop_size 360 --display_id -1 I also want to report that I had to edit the script unalligned_dataset.py since the Added letter prevented the script from finding the data

# self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')  
self.dir_A = os.path.join(opt.dataroot, opt.phase) 

# self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')  
self.dir_B = os.path.join(opt.dataroot, opt.phase)

finally my question is that my dataset is formed of complex objects I applied a set of masks and I need the model to remove those masks, for example I added a the some star object to the src image, my end goal is to use the network to remove that star from a new image that contains a different object but the same mask(in this example star). what is the recommended training epochs to achieve that goal?

maryam95hallal commented 1 year ago

I solved the issue if any needs help in that case contact me.

nik13 commented 1 year ago

Can you post the solution here? Thanks

maryam95hallal commented 1 year ago

@nik13 sorry for the late reply I moved on to work on other projects. If you still need the solution post your email and I will share a jupyter notebook of the process

here is a the solution !python train.py --model pix2pix --dataroot /content/gdrive/MyDrive/projects/pix2pix_official/dataset/rectangle_640_watermark_v1/dataset --name watermark_remover_pix2pix_rectangle_datasetv1 --direction AtoB --save_epoch_freq 50 --n_epochs 800 --n_epochs_decay 800 --preprocess crop --crop_size 256 --continue_train --epoch_count 1250 --display_id -1

During training, you can train do augmentation by resizing and random cropping. Since the network is fully convolutional, this means that you can test on any size of the image, and it won't be affected. To elaborate, assume that you have an image of 800x600, you can train on square patches of 256x26, but you set the --resize_or_crop none, during testing, so the full image (original resolution) will be translated normally based on your saved weights/model