generate_all.py

"""General-purpose training script for image-to-image translation.

This script works for various models (with option '--model': e.g., pix2pix, cyclegan, colorization) and
different datasets (with option '--dataset_mode': e.g., aligned, unaligned, single, colorization).
You need to specify the dataset ('--dataroot'), experiment name ('--name'), and model ('--model').

It first creates model, dataset, and visualizer given the option.
It then does standard network training. During the training, it also visualize/save the images, print/save the loss plot, and save models.
The script supports continue/resume training. Use '--continue_train' to resume your previous training.

Example:
    Train a CycleGAN model:
        python train.py --dataroot ./datasets/maps --name maps_cyclegan --model cycle_gan
    Train a pix2pix model:
        python train.py --dataroot ./datasets/facades --name facades_pix2pix --model pix2pix --direction BtoA

See options/base_options.py and options/train_options.py for more training options.
See training and test tips at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/tips.md
See frequently asked questions at: https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix/blob/master/docs/qa.md
"""
import time
from options.test_options import TestOptions
from datasets import create_dataset
from models import create_model
import os, torch, shutil
from tqdm import tqdm

if __name__ == '__main__':
    opt = TestOptions().parse()   # get training options
    if opt.square:
        opt.crop_size = (opt.crop_size, opt.crop_size)
    else:
        opt.crop_size = (opt.crop_size, max(1,int(opt.crop_size*1.0/256*176)))
    print("crop_size:", opt.crop_size)
    dataset = create_dataset(opt)  # create a dataset given opt.dataset_mode and other options
    dataset_size = len(dataset)    # get the number of images in the dataset.
    
    print('The number of training images = %d' % dataset_size)

    model = create_model(opt)      # create a model given opt.model and other options
    
    model.setup(opt)               # regular setup: load and print networks; create schedulers
    total_iters = 0                # the total number of training iterations
    model.eval()
    
    generate_out_dir = os.path.join(opt.eval_output_dir + "_%s"%opt.epoch)
    print("generate images at %s" % generate_out_dir)
    os.mkdir(generate_out_dir)
    model.isTrain = False
    # generate
    count = 0
    for i, data in tqdm(enumerate(dataset), "generating for test split"):  # inner loop within one epoch
        with torch.no_grad():
            model.set_input(data)         # unpack data from dataset and apply preprocessing
            model.forward()
            count = model.save_batch(generate_out_dir, count)