style.py

import numpy as np
import torch
import os
import argparse
import time
import collections

from torch.autograd import Variable
from torch.optim import Adam
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision import transforms

import utils
from network import ImageTransformNet, ImageTransformNet_dpws
from vgg import Vgg16

# Global Variables
IMAGE_SIZE = 256
BATCH_SIZE = 4
LEARNING_RATE = 1e-3
EPOCHS = 2

# STYLE_WEIGHT = 9.2e3
# STYLE_WEIGHT = 8e4
STYLE_WEIGHT = 7e4
# STYLE_WEIGHT = 0
# CONTENT_WEIGHT = 1e-2
# CONTENT_WEIGHT = 0.15
CONTENT_WEIGHT = 0.1
L1_WEIGHT = 1

def train(args):
    # GPU enabling
    if (args.gpu != None):
        use_cuda = True
        dtype = torch.cuda.FloatTensor
        torch.cuda.set_device(args.gpu)
        print("Current device: %d" %torch.cuda.current_device())

    # visualization of training controlled by flag
    visualize = (args.visualize != None)
    if (visualize):
        img_transform_512 = transforms.Compose([
            transforms.Scale(512),                  # scale shortest side to image_size
            # transforms.CenterCrop(512),           # crop center image_size out
            transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
            utils.normalize_tensor_transform()      # normalize with ImageNet values
        ])


        testImage_maine = utils.load_image(args.test_image)
        testImage_maine = img_transform_512(testImage_maine)
        testImage_maine = Variable(testImage_maine.repeat(1, 1, 1, 1), requires_grad=False).type(dtype)
        test_name = os.path.split(args.test_image)[-1].split('.')[0]

    # define network
    image_transformer_dpws = ImageTransformNet_dpws().type(dtype)
    # paras = [image_transformer_dpws.parameters()]
    optimizer = Adam(image_transformer_dpws.parameters(), LEARNING_RATE) 

    loss_mse = torch.nn.MSELoss()
    loss_l1 = torch.nn.L1Loss() 
   
    vgg = Vgg16().type(dtype)
    image_transformer = ImageTransformNet().type(dtype)
    image_transformer.load_state_dict(torch.load(args.load_path))

    # get training dataset
    dataset_transform = transforms.Compose([
        transforms.Scale(IMAGE_SIZE),           # scale shortest side to image_size
        transforms.CenterCrop(IMAGE_SIZE),      # crop center image_size out
        transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
        utils.normalize_tensor_transform()      # normalize with ImageNet values
    ])
    train_dataset = datasets.ImageFolder(args.dataset, dataset_transform)
    train_loader = DataLoader(train_dataset, batch_size = BATCH_SIZE)

    # style image
    style_transform = transforms.Compose([
        transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
        utils.normalize_tensor_transform()      # normalize with ImageNet values
    ])
    style = utils.load_image(args.style_image)
    style = style_transform(style)
    style = Variable(style.repeat(BATCH_SIZE, 1, 1, 1)).type(dtype)
    style_name = os.path.split(args.style_image)[-1].split('.')[0]
    
    # calculate gram matrices for style feature layer maps we care about
    style_features = vgg(style)
    style_gram = [utils.gram(fmap) for fmap in style_features]

    for e in range(EPOCHS):

        # track values for...
        img_count = 0
        aggregate_style_loss = 0.0
        aggregate_content_loss = 0.0
        aggregate_l1_loss = 0.0
        # aggregate_tv_loss = 0.0

        # train network
        image_transformer_dpws.train()
        for batch_num, (x, label) in enumerate(train_loader):
            img_batch_read = len(x)
            img_count += img_batch_read

            # zero out gradients
            optimizer.zero_grad()

            # input batch to transformer network

            x = Variable(x).type(dtype)
            y_hat = image_transformer_dpws(x)
            y_label = image_transformer(x)

            # get vgg features
            y_c_features = vgg(x)
            y_hat_features = vgg(y_hat)

            # calculate style loss
            y_hat_gram = [utils.gram(fmap) for fmap in y_hat_features]
            style_loss = 0.0
            
            for j in range(4):
                style_loss += loss_mse(y_hat_gram[j], style_gram[j][:img_batch_read])
            style_loss = STYLE_WEIGHT*style_loss
            aggregate_style_loss += style_loss.data.item()

            # calculate content loss (h_relu_2_2)
            recon = y_c_features[1]      
            recon_hat = y_hat_features[1]
            content_loss = CONTENT_WEIGHT*loss_mse(recon_hat, recon)
            aggregate_content_loss += content_loss.data.item()

            # calculate l1 loss
            l1_loss = L1_WEIGHT*loss_mse(y_hat, y_label)
            aggregate_l1_loss += l1_loss.data.item()

            # total loss
            # total_loss = style_loss + content_loss + tv_loss + l1_loss + dis_loss
            total_loss = style_loss + l1_loss + content_loss

            # backprop
            total_loss.backward()
            optimizer.step()

            # print out status message
            if ((batch_num + 1) % 100 == 0):
                status = "{}  Epoch {}:  [{}/{}]  Batch:[{}]  agg_style: {:.6f}  agg_l1: {:.6f}  agg_content: {:.6f} ".format(
                                time.ctime(), e + 1, img_count, len(train_dataset), batch_num+1,
                                aggregate_style_loss/(batch_num+1.0), aggregate_l1_loss/(batch_num+1.0), aggregate_content_loss/(batch_num+1.0)
                                )
                print(status)

            if ((batch_num + 1) % 5000 == 0) and (visualize):
                image_transformer_dpws.eval()

                if not os.path.exists("visualization"):
                    os.makedirs("visualization")

                outputTestImage_maine = image_transformer_dpws(testImage_maine)
    
                test_path = "visualization/%s/%s%d_%05d.jpg" %(style_name, test_name, e+1, batch_num+1)
                utils.save_image(test_path, outputTestImage_maine.data[0].cpu())

                print("images saved")
                image_transformer_dpws.train()

    # save model
    image_transformer_dpws.eval()

    if use_cuda:
        image_transformer_dpws.cpu()

    if not os.path.exists("models"):
        os.makedirs("models")
    filename = "models/%s.model" %style_name
    torch.save(image_transformer_dpws.state_dict(), filename)
    
    if use_cuda:
        image_transformer_dpws.cuda()

def style_transfer(args):
    # GPU enabling
    if (args.gpu != None):
        use_cuda = True
        dtype = torch.cuda.FloatTensor
        torch.cuda.set_device(args.gpu)
        print("Current device: %d" %torch.cuda.current_device())
    else : 
        dtype = torch.FloatTensor
        
    # content image
    img_transform_512 = transforms.Compose([
            # transforms.Scale(512),                  # scale shortest side to image_size
            transforms.Resize(512),                  # scale shortest side to image_size
            # transforms.CenterCrop(512),             # crop center image_size out
            transforms.ToTensor(),                  # turn image from [0-255] to [0-1]
            utils.normalize_tensor_transform()      # normalize with ImageNet values
    ])

    content = utils.load_image(args.source)
    content = img_transform_512(content)
    content = content.unsqueeze(0)
    # content = Variable(content).type(dtype)
    content = Variable(content.repeat(1, 1, 1, 1), requires_grad=False).type(dtype)

    # load style model
    checkpoint_lw = torch.load(args.model_path)

    style_model = ImageTransformNet_dpws().type(dtype)
    style_model.load_state_dict((checkpoint_lw))

    # process input image
    stylized = style_model(content).cpu()
    utils.save_image(args.output, stylized.data[0])


def main():
    parser = argparse.ArgumentParser(description='style transfer in pytorch')
    subparsers = parser.add_subparsers(title="subcommands", dest="subcommand")

    train_parser = subparsers.add_parser("train", help="train a model to do style transfer")
    train_parser.add_argument("--style_image", type=str, required=True, help="path to a style image to train with")
    train_parser.add_argument("--test_image", type=str, required=True, help="path to a test image to test with")
    train_parser.add_argument("--dataset", type=str, required=True, help="path to a dataset")
    train_parser.add_argument("--gpu", type=int, default=None, help="ID of GPU to be used")
    train_parser.add_argument("--visualize", type=int, default=None, help="Set to 1 if you want to visualize training")

    style_parser = subparsers.add_parser("transfer", help="do style transfer with a trained model")
    style_parser.add_argument("--model_path", type=str, required=True, help="path to a pretrained model for a style image")
    style_parser.add_argument("--source", type=str, required=True, help="path to source image")
    style_parser.add_argument("--output", type=str, required=True, help="file name for stylized output image")
    style_parser.add_argument("--gpu", type=int, default=None, help="ID of GPU to be used")

    args = parser.parse_args()

    # command
    if (args.subcommand == "train"):
        print("Training!")
        train(args)
    elif (args.subcommand == "transfer"):
        print("Style transfering!")
        style_transfer(args)
    else:
        print("invalid command")

if __name__ == '__main__':
    main()