<repo_name>organizations-loops/SR_2020<file_sep>/main.py
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
import os
import argparse
from solver import Solver
from data_loader import get_loader
from torch.backends import cudnn


def str2bool(v):
    return v.lower() in ('true')

def main(config):
    # For fast training.
    cudnn.benchmark = True

    # Create directories if not exist.
    if not os.path.exists(config.log_dir):
        os.makedirs(config.log_dir)
    if not os.path.exists(config.model_save_dir):
        os.makedirs(config.model_save_dir)
    if not os.path.exists(config.sample_dir):
        os.makedirs(config.sample_dir)
    if not os.path.exists(config.result_dir):
        os.makedirs(config.result_dir)

    # Data loader.
    celeba_loader = get_loader(config.celeba_image_dir, config.attr_path, config.selected_attrs,
                               config.image_size, config.magnification, config.batch_size, 
                               config.dataset, config.mode, config.num_workers)

    # Solver for training and testing our networks.
    solver = Solver(celeba_loader, config)
    if config.mode == 'train':
        solver.train()
    elif config.mode == 'test':
        solver.test()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    # Model configuration.
    parser.add_argument('--c_dim', type=int, default=4, help='dimension of domain labels (1st dataset)')
    parser.add_argument('--image_size', type=int, default=128, help='output image resolution')
    parser.add_argument('--magnification', type=int, default=4, choices=[2, 4, 8, 16], help='image magnification scale')
    parser.add_argument('--tg_conv_dim', type=int, default=64, help='number of conv filters in the first layer of T_G')
    parser.add_argument('--td_conv_dim', type=int, default=64, help='number of conv filters in the first layer of T_D')
    parser.add_argument('--tg_repeat_num', type=int, default=6, help='number of residual blocks in TG')
    parser.add_argument('--td_repeat_num', type=int, default=6, help='number of strided conv layers in TD')
    parser.add_argument('--lambda_cls', type=float, default=1, help='weight for domain classification loss')
    parser.add_argument('--lambda_rec', type=float, default=10, help='weight for reconstruction loss')
    parser.add_argument('--lambda_gp', type=float, default=10, help='weight for gradient penalty')

    # Training configuration.
    parser.add_argument('--dataset', type=str, default='CelebA', choices=['CelebA'], help='dataset is CelebA')
    parser.add_argument('--batch_size', type=int, default=16, help='mini-batch size')
    parser.add_argument('--num_iters', type=int, default=300000, help='number of total iterations for training T_D')
    parser.add_argument('--num_iters_decay', type=int, default=100000, help='number of iterations for decaying lr')
    parser.add_argument('--tg_lr', type=float, default=1e-4, help='learning rate for T_G')
    parser.add_argument('--td_lr', type=float, default=1e-4, help='learning rate for T_D')
    parser.add_argument('--eg_lr', type=float, default=1e-4, help='learning rate for E_G')
    parser.add_argument('--edc_lr', type=float, default=1e-4, help='learning rate for E_Dc')
    parser.add_argument('--edt_lr', type=float, default=1e-4, help='learning rate for E_Dt')
    parser.add_argument('--tv_weight', type=float, default=10, help='tv_loss weight')
    parser.add_argument('--n_critic', type=int, default=5, help='number of T_D updates per each T_G update')
    parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for Adam optimizer')
    parser.add_argument('--beta2', type=float, default=0.999, help='beta2 for Adam optimizer')
    parser.add_argument('--resume_iters', type=int, default=None, help='resume training from this step')
    parser.add_argument('--selected_attrs', '--list', nargs='+', help='selected attributes for the CelebA dataset', 
                        default=['Male', 'Mustache', 'Big_Nose', 'Mouth_Slightly_Open'])
    # default=['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Male', 'Young']
    # default=['Male', 'Young', 'Bags_Under_Eyes', 'Heavy_Makeup', 'No_Beard', '5_o_Clock_Shadow', 'Mustache', 'Pointy_Nose', 'Big_Nose', 'Blurry', 'Narrow_Eyes', 'Eyeglasses', 'Smiling', 'Mouth_Slightly_Open']

    # Test configuration.
    parser.add_argument('--test_iters', type=int, default=300000, help='test model from this step')

    # Miscellaneous.
    parser.add_argument('--num_workers', type=int, default=1)
    parser.add_argument('--mode', type=str, default='train', choices=['train', 'test'])
    parser.add_argument('--use_tensorboard', type=str2bool, default=True)

    # Directories.
    parser.add_argument('--celeba_image_dir', type=str, default='/media/limy/新加卷1/Datasets/Faces/CelebA/Celeba-HQ/celeba-128')
    parser.add_argument('--attr_path', type=str, default='/media/limy/新加卷1/Datasets/Faces/CelebA/Anno/list_attr_celebahq.txt')
    parser.add_argument('--log_dir', type=str, default='stargan/logs')
    parser.add_argument('--model_save_dir', type=str, default='stargan/models')
    parser.add_argument('--sample_dir', type=str, default='stargan/samples')
    parser.add_argument('--result_dir', type=str, default='stargan/results')

    # Step size.
    parser.add_argument('--log_step', type=int, default=10)
    parser.add_argument('--sample_step', type=int, default=1000)
    parser.add_argument('--model_save_step', type=int, default=10000)
    parser.add_argument('--lr_update_step', type=int, default=1000)

    config = parser.parse_args()
    print(config)
    main(config)
<file_sep>/solver.py
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
from model import Transform_Generator
from model import Transform_Discriminator
from model import Enhance_Generator
from model import Enhance_DiscriminatorC, Enhance_DiscriminatorT
from model import TVLoss
from model import GaussianBlur
from model import GrayLayer
from model import VGG
from torch.autograd import Variable
from torchvision.utils import save_image
from torchvision import transforms as T
import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
import numpy as np
import os
import time
import datetime



class Solver(object):
    """Solver for training and testing StarGAN."""

    def __init__(self, celeba_loader, config):
        """Initialize configurations."""

        # Data loader.
        self.celeba_loader = celeba_loader

        # Model configurations.
        self.c_dim = config.c_dim

        self.image_size = config.image_size
        self.magnification = config.magnification
        ## Transform Network configurations.
        self.tg_conv_dim = config.tg_conv_dim
        self.td_conv_dim = config.td_conv_dim
        self.tg_repeat_num = config.tg_repeat_num
        self.td_repeat_num = config.td_repeat_num
        self.lambda_cls = config.lambda_cls
        self.lambda_rec = config.lambda_rec
        self.lambda_gp = config.lambda_gp
        ## Transform Network configurations.
        
        
        # Training configurations.
        self.dataset = config.dataset
        self.batch_size = config.batch_size
        self.num_iters = config.num_iters
        self.num_iters_decay = config.num_iters_decay
        self.tg_lr = config.tg_lr
        self.td_lr = config.td_lr
        self.eg_lr = config.eg_lr
        self.edc_lr = config.edc_lr
        self.edt_lr = config.edt_lr
        self.n_critic = config.n_critic
        self.beta1 = config.beta1
        self.beta2 = config.beta2
        self.resume_iters = config.resume_iters
        self.selected_attrs = config.selected_attrs

        # Test configurations.
        self.test_iters = config.test_iters

        # Miscellaneous.
        self.use_tensorboard = config.use_tensorboard
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

        # Directories.
        self.log_dir = config.log_dir
        self.sample_dir = config.sample_dir
        self.model_save_dir = config.model_save_dir
        self.result_dir = config.result_dir

        # Step size.
        self.log_step = config.log_step
        self.sample_step = config.sample_step
        self.model_save_step = config.model_save_step
        self.lr_update_step = config.lr_update_step

        # TV loss
        self.tv_criterion = TVLoss(config.tv_weight)
        # Color loss
        self.color_criterion = nn.CrossEntropyLoss()
        # Texture loss
        self.texture_criterion = nn.CrossEntropyLoss()
        # identity loss
        self.identity_criterion = nn.L1Loss()
        # content loss
        self.content_criterion = nn.MSELoss()
        # reconstruction loss
        self.rec_criterion = nn.MSELoss()

        # Enhancement Operations
        self.blur = GaussianBlur()
        self.gray = GrayLayer()
        self.vgg = VGG()

        # Build the model and tensorboard.
        self.build_model()
        if self.use_tensorboard:
            self.build_tensorboard()

    def build_model(self):
        """Create generators and discriminators: Transform_Generator(T_G), Transform_Discriminator(T_D), 
        Enhance_Generator(E_G), Enhance_Discriminator_color(E_Dc), Enhance_Discriminator_texture(E_Dt)."""

        self.T_G = Transform_Generator(self.tg_conv_dim, self.c_dim, self.tg_repeat_num)
        self.T_D = Transform_Discriminator(self.image_size, self.td_conv_dim, self.c_dim, self.td_repeat_num) 
        self.E_G = Enhance_Generator()
        
        self.tg_optimizer = torch.optim.Adam(self.T_G.parameters(), self.tg_lr, [self.beta1, self.beta2])
        self.td_optimizer = torch.optim.Adam(self.T_D.parameters(), self.td_lr, [self.beta1, self.beta2])
        self.eg_optimizer = torch.optim.Adam(self.E_G.parameters(), self.eg_lr, [self.beta1, self.beta2])
        
        self.print_network(self.T_G, 'T_G')
        self.print_network(self.T_D, 'T_D')
        self.print_network(self.E_G, 'E_G')
        
        self.T_G.to(self.device)
        self.T_D.to(self.device)
        self.E_G.to(self.device)

    def print_network(self, model, name):
        """Print out the network information."""
        num_params = 0
        for p in model.parameters():
            num_params += p.numel()
        print(model)
        print(name)
        print("The number of parameters: {}".format(num_params))

    def restore_model(self, resume_iters):
        """Restore the trained generator and discriminator."""
        print('Loading the trained models from step {}...'.format(resume_iters))
        T_G_path = os.path.join(self.model_save_dir, '{}-T_G.ckpt'.format(resume_iters))
        T_D_path = os.path.join(self.model_save_dir, '{}-T_D.ckpt'.format(resume_iters))
        E_G_path = os.path.join(self.model_save_dir, '{}-E_G.ckpt'.format(resume_iters))

        self.T_G.load_state_dict(torch.load(T_G_path, map_location=lambda storage, loc: storage))
        self.T_D.load_state_dict(torch.load(T_D_path, map_location=lambda storage, loc: storage))
        self.E_G.load_state_dict(torch.load(E_G_path, map_location=lambda storage, loc: storage))


    def build_tensorboard(self):
        """Build a tensorboard logger."""
        from logger import Logger
        self.logger = Logger(self.log_dir)

    def update_lr(self, tg_lr, td_lr, eg_lr, edc_lr, edt_lr):
        """Decay learning rates of the generator and discriminator."""
        for param_group in self.tg_optimizer.param_groups:
            param_group['lr'] = tg_lr
        for param_group in self.td_optimizer.param_groups:
            param_group['lr'] = td_lr
        for param_group in self.eg_optimizer.param_groups:
            param_group['lr'] = eg_lr


    def reset_grad(self):
        """Reset the gradient buffers."""
        self.tg_optimizer.zero_grad()
        self.td_optimizer.zero_grad()
        self.eg_optimizer.zero_grad()


    def denorm(self, x):
        """Convert the range from [-1, 1] to [0, 1]."""
        out = (x + 1) / 2
        return out.clamp_(0, 1)

    def gradient_penalty(self, y, x):
        """Compute gradient penalty: (L2_norm(dy/dx) - 1)**2."""
        weight = torch.ones(y.size()).to(self.device)
        dydx = torch.autograd.grad(outputs=y,
                                   inputs=x,
                                   grad_outputs=weight,
                                   retain_graph=True,
                                   create_graph=True,
                                   only_inputs=True)[0]

        dydx = dydx.view(dydx.size(0), -1)
        dydx_l2norm = torch.sqrt(torch.sum(dydx**2, dim=1))
        return torch.mean((dydx_l2norm-1)**2)

    def label2onehot(self, labels, dim):
        """Convert label indices to one-hot vectors."""
        batch_size = labels.size(0)
        out = torch.zeros(batch_size, dim)
        out[np.arange(batch_size), labels.long()] = 1
        return out

    def create_labels(self, c_org, c_dim=5, dataset='CelebA', selected_attrs=None):
        """Generate target domain labels for debugging and testing."""
        # Get hair color indices.
        if dataset == 'CelebA':
            hair_color_indices = []
            for i, attr_name in enumerate(selected_attrs):
                if attr_name in ['Black_Hair', 'Blond_Hair', 'Brown_Hair', 'Gray_Hair']:
                    hair_color_indices.append(i)

        c_trg_list = []
        for i in range(c_dim):
            if dataset == 'CelebA':
                c_trg = c_org.clone()
                if i in hair_color_indices:  # Set one hair color to 1 and the rest to 0.
                    c_trg[:, i] = 1
                    for j in hair_color_indices:
                        if j != i:
                            c_trg[:, j] = 0
                else:
                    c_trg[:, i] = (c_trg[:, i] == 0)  # Reverse attribute value.

            c_trg_list.append(c_trg.to(self.device))
        return c_trg_list

    def classification_loss(self, logit, target):
        """Compute binary or softmax cross entropy loss."""
        return F.binary_cross_entropy_with_logits(logit, target, size_average=False) / logit.size(0)

    def train(self):
        """Train StarGAN within a single dataset."""
        # Set data loader.
        data_loader = self.celeba_loader

        # Fetch fixed inputs for debugging.
        data_iter = iter(data_loader)
        x_H_fixed, x_Lm_fixed, x_L_fixed, y_org = next(data_iter)
        x_H_fixed = x_H_fixed.to(self.device)
        x_Lm_fixed = x_Lm_fixed.to(self.device)
        x_L_fixed = x_L_fixed.to(self.device)
        y_trg_list = self.create_labels(y_org, self.c_dim, self.dataset, self.selected_attrs)

        # Learning rate cache for decaying.
        tg_lr = self.tg_lr
        td_lr = self.td_lr
        eg_lr = self.eg_lr

        # Start training from scratch or resume training.
        start_iters = 0
        if self.resume_iters:
            start_iters = self.resume_iters
            self.restore_model(self.resume_iters)

        # Start training.
        print('Start training...')
        start_time = time.time()
        for i in range(start_iters, self.num_iters):

            # =================================================================================== #
            #                             1. Preprocess input data                                #
            # =================================================================================== #

            # Fetch real images and labels.
            try:
                x_H, x_Lm, x_L, label_org = next(data_iter)
            except:
                data_iter = iter(data_loader)
                x_H, x_Lm, x_L, label_org = next(data_iter)

            # Generate target domain labels randomly.
            rand_idx = torch.randperm(label_org.size(0))
            label_trg = label_org[rand_idx]
            # lab

            y_org = label_org.clone()
            y_trg = label_trg.clone()

            x_H = x_H.to(self.device)                 # HR images
            x_Lm = x_Lm.to(self.device)               # LR images
            x_L = x_L.to(self.device)                 # Input images.
            y_org = y_org.to(self.device)             # Original domain labels.
            y_trg = y_trg.to(self.device)             # Target domain labels.
            label_org = label_org.to(self.device)     # Labels for computing classification loss.
            label_trg = label_trg.to(self.device)     # Labels for computing classification loss.

            
            # =================================================================================== #
            #                             2. Generators and Discriminators                        #
            # =================================================================================== #
            """
            # Transform_Generator
            x_T = self.T_G(x_L, y_trg)
            x_T_TG = self.T_G(x_T, y_org)
            
            # Transform_Discriminator & Classifier
            x_T_TD, x_T_TC = self.T_D(x_T)
            x_H_TD, x_H_TC = self.T_D(x_H)
            
            # Enhancement_Generator
            x_E = self.E_G(x_T)
            
            # Enhancement_VGG for identity loss
            x_E_vgg = self.vgg(x_E)
            x_H_vgg = self.vgg(x_H)
            
            # Enhancement_Discriminator_Color for color loss
            x_E_blur = self.blur(x_E)
            x_H_blur = self.blur(x_H)
            x_E_blur_EDc = self.E_Dc(x_E_blur)
            x_H_blur_EDc = self.E_Dc(x_H_blur)

            # Enhancement_Discriminator_Texture for texture loss
            x_E_gray = self.gray(x_E)
            x_H_gray = self.gray(x_H)
            x_E_gray_EDt = self.E_Dt(x_E_gray)
            x_H_gray_EDt = self.E_Dt(x_H_gray)
            """

            # =================================================================================== #
            #                             2. Train the discriminator                              #
            # =================================================================================== #
            
            # Discriminators
            ## Transform Network
            x_H_TD, x_H_TC = self.T_D(x_H)
            x_T = self.T_G(x_L, y_trg)
            x_T_TD, x_T_TC = self.T_D(x_T.detach())
            ### adv_loss
            td_loss_real = - torch.mean(x_H_TD)
            td_loss_fake = torch.mean(x_T_TD)
            td_loss_adv = td_loss_real + td_loss_fake
            ### att_loss
            td_loss_att = self.classification_loss(x_H_TC, label_org)
            ### gp_loss
            alpha = torch.rand(x_H.size(0), 1, 1, 1).to(self.device)
            x_hat = (alpha * x_H.data + (1 - alpha) * x_T.data).requires_grad_(True)
            x_hat_TD, _ = self.T_D(x_hat)
            td_loss_gp = self.gradient_penalty(x_hat_TD, x_hat)
            ### T_D Total Loss
            td_loss = td_loss_adv + td_loss_att + td_loss_gp * 10
            ## Enhancement Network
            x_E = self.E_G(x_T)

            # Discriminators Loss
            d_loss = td_loss
            
            self.reset_grad()
            d_loss.backward()
            self.td_optimizer.step()

            
            # Logging.
            loss = {}
            loss['T_D/loss_TD'] = td_loss.item()
            loss['T_D/loss_adv'] = td_loss_adv.item()
            loss['T_D/loss_att'] = td_loss_att.item()
            loss['T_D/loss_gp'] = td_loss_gp.item()

            
            # =================================================================================== #
            #                               3. Train the generators                               #
            # =================================================================================== #
            
            if (i+1) % self.n_critic == 0:
                
                # Generators
                ## Transform Network
                x_T = self.T_G(x_L, y_trg)
                x_T_TD, x_T_TC = self.T_D(x_T)
                x_T_TG = self.T_G(x_T, y_org)
                ### adv_loss
                tg_loss_fake = - torch.mean(x_T_TD)
                tg_loss_adv = tg_loss_fake
                ### att_loss
                tg_loss_att = self.classification_loss(x_T_TC, label_trg)
                ### rec_loss
                # tg_loss_rec = torch.mean(torch.abs(x_H - x_T_TG))
                tg_loss_rec = self.rec_criterion(x_T_TG, x_H)

                ### T_G Total Loss
                tg_loss = tg_loss_adv + tg_loss_att + tg_loss_rec * 10 + tg_loss_tv

                ### identity_loss
                _, c1, h1, w1 = x_T.size()
                chw1 = c1 * h1 * w1
                eg_loss_identity = 1.0/chw1 * self.content_criterion(x_E_vgg, x_H_vgg)

                ### content_loss
                eg_loss_content = self.content_criterion(x_E, x_H)
                ### E_G Total Loss
                eg_loss = eg_loss_identity + eg_loss_content
                # Generators Loss
                g_loss = tg_loss + eg_loss
                
                self.reset_grad()
                g_loss.backward()
                self.tg_optimizer.step()
                self.eg_optimizer.step()
                
                # Logging.
                loss = {}
                loss['T_G/loss_TG'] = tg_loss.item()
                loss['T_G/loss_adv'] = tg_loss_adv.item()
                loss['T_G/loss_att'] = tg_loss_att.item()
                loss['T_G/loss_rec'] = tg_loss_rec.item()
                loss['E_G/loss_EG'] = eg_loss.item()
                loss['E_G/loss_identity'] = eg_loss_identity.item()
                loss['E_G/loss_content'] = eg_loss_content.item()
            
            # =================================================================================== #
            #                                 4. Miscellaneous                                    #
            # =================================================================================== #

            # Print out training information.
            if (i+1) % self.log_step == 0:
                et = time.time() - start_time
                et = str(datetime.timedelta(seconds=et))[:-7]
                log = "Elapsed [{}], Iteration [{}/{}]".format(et, i+1, self.num_iters)
                for tag, value in loss.items():
                    log += ", {}: {:.4f}".format(tag, value)
                print(log)

                if self.use_tensorboard:
                    for tag, value in loss.items():
                        self.logger.scalar_summary(tag, value, i+1)

            # Translate fixed images for debugging.
            if (i+1) % self.sample_step == 0:
                with torch.no_grad():
                    x_output_list = [x_H_fixed, x_Lm_fixed, x_L_fixed]
                    x_output_list.append(self.T_G(x_L_fixed, y_org))
                    x_output_list.append(self.E_G(self.T_G(x_L_fixed, y_org)))
                    for y_trg in y_trg_list:
                        x_output_list.append(self.T_G(x_L_fixed, y_trg))
                        x_output_list.append(self.E_G(self.T_G(x_L_fixed, y_trg)))
                    x_concat = torch.cat(x_output_list, dim=3)
                    sample_path = os.path.join(self.sample_dir, '{}-images.jpg'.format(i+1))
                    save_image(self.denorm(x_concat.data.cpu()), sample_path, nrow=1, padding=0)
                    print('Saved real and fake images into {}...'.format(sample_path))

            # Save model checkpoints.
            if (i+1) % self.model_save_step == 0:
                T_G_path = os.path.join(self.model_save_dir, '{}-T_G.ckpt'.format(i+1))
                T_D_path = os.path.join(self.model_save_dir, '{}-T_D.ckpt'.format(i+1))
                E_G_path = os.path.join(self.model_save_dir, '{}-E_G.ckpt'.format(i+1))
                torch.save(self.T_G.state_dict(), T_G_path)
                torch.save(self.T_D.state_dict(), T_D_path)
                torch.save(self.E_G.state_dict(), E_G_path)
                print('Saved model checkpoints into {}...'.format(self.model_save_dir))

            # Decay learning rates.
            if (i+1) % self.lr_update_step == 0 and (i+1) > (self.num_iters - self.num_iters_decay):
                tg_lr -= (self.tg_lr / float(self.num_iters_decay))
                td_lr -= (self.td_lr / float(self.num_iters_decay))
                eg_lr -= (self.eg_lr / float(self.num_iters_decay))
                self.update_lr(tg_lr, td_lr, eg_lr)
                print ('Decayed learning rates, tg_lr: {}, td_lr: {}, eg_lr: {}, edc_lr: {}, edt_lr: {}.'.format(tg_lr, td_lr, eg_lr, edc_lr, edt_lr))

    def test(self):
        """Translate images using StarGAN trained on a single dataset."""
        # Load the trained generator.
        self.restore_model(self.test_iters)
        
        # Set data loader.
        data_loader = self.celeba_loader

        with torch.no_grad():
            for i, (x_H, x_Lm, x_L, y_org) in enumerate(data_loader):

                # Prepare input images and target domain labels.
                x_H = x_H.to(self.device)
                x_Lm = x_Lm.to(self.device)
                x_L = x_L.to(self.device)
                y_org = y_org.to(self.device)
                y_trg_list = self.create_labels(y_org, self.c_dim, self.dataset, self.selected_attrs)

                # Translate images.
                x_output_list = [x_H, x_Lm, x_L]
                x_output_list.append(self.T_G(x_L, y_org))
                x_output_list.append(self.E_G(self.T_G(x_L, y_org)))
                for y_trg in y_trg_list:
                    x_output_list.append(self.T_G(x_L, y_trg))
                    x_output_list.append(self.E_G(self.T_G(x_L, y_trg)))

                # Save the translated images.
                x_concat = torch.cat(x_output_list, dim=3)
                result_path = os.path.join(self.result_dir, '{}-images.jpg'.format(i+1))
                save_image(self.denorm(x_concat.data.cpu()), result_path, nrow=1, padding=0)
                print('Saved real and fake images into {}...'.format(result_path))

<file_sep>/data_loader.py
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
from torch.utils import data
from torchvision import transforms as T
from torchvision.datasets import ImageFolder
from PIL import Image
import torch
import os
import random


class CelebA(data.Dataset):
    """Dataset class for the CelebA dataset."""

    def __init__(self, image_dir, attr_path, selected_attrs, transform, mode):
        """Initialize and preprocess the CelebA dataset."""
        self.image_dir = image_dir
        self.attr_path = attr_path
        self.selected_attrs = selected_attrs
        self.transform1 = transform[0]
        self.transform2 = transform[1]
        self.transform3 = transform[2]
        self.mode = mode
        self.train_dataset = []
        self.test_dataset = []
        self.attr2idx = {}
        self.idx2attr = {}
        self.preprocess()

        if mode == 'train':
            self.num_images = len(self.train_dataset)
        else:
            self.num_images = len(self.test_dataset)

    def preprocess(self):
        """Preprocess the CelebA attribute file."""
        lines = [line.rstrip() for line in open(self.attr_path, 'r')]
        all_attr_names = lines[1].split()
        for i, attr_name in enumerate(all_attr_names):
            self.attr2idx[attr_name] = i
            self.idx2attr[i] = attr_name

        lines = lines[2:]
        random.seed(1234)
        random.shuffle(lines)
        for i, line in enumerate(lines):
            split = line.split()
            filename = split[0]
            values = split[1:]

            label = []
            for attr_name in self.selected_attrs:
                idx = self.attr2idx[attr_name]
                label.append(values[idx] == '1')

            if (i+1) < 30000:
                self.test_dataset.append([filename, label])
            else:
                self.train_dataset.append([filename, label])

        print('Finished preprocessing the CelebA dataset...')

    def __getitem__(self, index):
        """Return three image and its corresponding attribute label."""
        dataset = self.train_dataset if self.mode == 'train' else self.test_dataset
        filename, label = dataset[index]
        image = Image.open(os.path.join(self.image_dir, filename))
        # Return: HR images, LR images, Input images, Selected attribute labels
        return self.transform1(image), self.transform2(image), self.transform3(image), torch.FloatTensor(label)

    def __len__(self):
        """Return the number of images."""
        return self.num_images


def get_loader(image_dir, attr_path, selected_attrs, image_size=128, magnification=4,
               batch_size=16, dataset='CelebA', mode='train', num_workers=1):
    """Build and return a data loader."""

    # Transform1, Ground truth HR images to Tensor
    transform1 = []
    transform1.append(T.ToTensor())
    transform1.append(T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
    transform1 = T.Compose(transform1)

    # Transform2, LR images to Tensor
    transform2 = []
    transform2.append(T.Resize(image_size // magnification, interpolation=Image.BICUBIC))
    transform2.append(T.Resize(image_size, interpolation=Image.NEAREST))
    transform2.append(T.ToTensor())
    transform2.append(T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
    transform2 = T.Compose(transform2)

    # Transform3, HR ->downsample-> LR ->upsample-> Input images to Tensor
    transform3 = []
    # if mode == 'train':
    #     transform3.append(T.RandomHorizontalFlip())
    transform3.append(T.Resize(image_size // magnification, interpolation=Image.BICUBIC))
    transform3.append(T.Resize(image_size, interpolation=Image.BICUBIC))
    transform3.append(T.ToTensor())
    transform3.append(T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)))
    transform3 = T.Compose(transform3)

    # Transform
    transform = [transform1, transform2, transform3]

    if dataset == 'CelebA':
        dataset = CelebA(image_dir, attr_path, selected_attrs, transform, mode)

    data_loader = data.DataLoader(dataset=dataset,
                                  batch_size=batch_size,
                                  shuffle=(mode=='train'),
                                  num_workers=num_workers)
    return data_loader
<file_sep>/model.py
#!/usr/bin/env python2
# -*- coding: utf-8 -*-
import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision.models import vgg19_bn
import numpy as np


class ResidualBlock(nn.Module):
    """Residual Block with instance normalization."""
    def __init__(self, dim_in, dim_out):
        super(ResidualBlock, self).__init__()
        self.main = nn.Sequential(
            nn.Conv2d(dim_in, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
            nn.InstanceNorm2d(dim_out, affine=True, track_running_stats=True),
            nn.ReLU(inplace=True),
            nn.Conv2d(dim_out, dim_out, kernel_size=3, stride=1, padding=1, bias=False),
            nn.InstanceNorm2d(dim_out, affine=True, track_running_stats=True))

    def forward(self, x):
        return x + self.main(x)


class Transform_Generator(nn.Module):
    """Generator network."""
    def __init__(self, conv_dim=64, c_dim=5, repeat_num=6):
        super(Transform_Generator, self).__init__()

        layers = []
        layers.append(nn.Conv2d(3+c_dim, conv_dim, kernel_size=7, stride=1, padding=3, bias=False))
        layers.append(nn.InstanceNorm2d(conv_dim, affine=True, track_running_stats=True))
        layers.append(nn.ReLU(inplace=True))

        # Down-sampling layers.
        curr_dim = conv_dim
        for i in range(2):
            layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim*2, affine=True, track_running_stats=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim * 2

        # Bottleneck layers.
        for i in range(repeat_num):
            layers.append(ResidualBlock(dim_in=curr_dim, dim_out=curr_dim))

        # Up-sampling layers.
        for i in range(2):
            layers.append(nn.ConvTranspose2d(curr_dim, curr_dim//2, kernel_size=4, stride=2, padding=1, bias=False))
            layers.append(nn.InstanceNorm2d(curr_dim//2, affine=True, track_running_stats=True))
            layers.append(nn.ReLU(inplace=True))
            curr_dim = curr_dim // 2

        layers.append(nn.Conv2d(curr_dim, 3, kernel_size=7, stride=1, padding=3, bias=False))
        layers.append(nn.Tanh())
        self.main = nn.Sequential(*layers)

    def forward(self, x, c):
        # Replicate spatially and concatenate domain information.
        c = c.view(c.size(0), c.size(1), 1, 1)
        c = c.repeat(1, 1, x.size(2), x.size(3))
        x = torch.cat([x, c], dim=1)
        return self.main(x)


class Transform_Discriminator(nn.Module):
    """Discriminator network with PatchGAN."""
    def __init__(self, image_size=128, conv_dim=64, c_dim=5, repeat_num=6):
        super(Transform_Discriminator, self).__init__()
        layers = []
        layers.append(nn.Conv2d(3, conv_dim, kernel_size=4, stride=2, padding=1))
        layers.append(nn.LeakyReLU(0.01))

        curr_dim = conv_dim
        for i in range(1, repeat_num):
            layers.append(nn.Conv2d(curr_dim, curr_dim*2, kernel_size=4, stride=2, padding=1))
            layers.append(nn.LeakyReLU(0.01))
            curr_dim = curr_dim * 2

        kernel_size = int(image_size / np.power(2, repeat_num))
        self.main = nn.Sequential(*layers)
        self.conv1 = nn.Conv2d(curr_dim, 1, kernel_size=3, stride=1, padding=1, bias=False)
        self.conv2 = nn.Conv2d(curr_dim, c_dim, kernel_size=kernel_size, bias=False)
        
    def forward(self, x):
        h = self.main(x)
        out_src = self.conv1(h)
        out_cls = self.conv2(h)
        return out_src, out_cls.view(out_cls.size(0), out_cls.size(1))


class TVLoss(nn.Module):

    def __init__(self, tv_weight):
        super(TVLoss, self).__init__()
        self.tv_weight = tv_weight

    def forward(self, x):

        batch_size = x.size()[0]
        h_x = x.size()[2]
        w_x = x.size()[3]
        count_h = self._tensor_size(x[:, :, 1:, :])
        count_w = self._tensor_size(x[:, :, :, 1:])
        h_tv = torch.pow((x[:, :, 1:, :]-x[:, :, :h_x-1, :]), 2).sum()
        w_tv = torch.pow((x[:, :, :, 1:]-x[:, :, :, :w_x-1]),2).sum()
        return self.tv_weight*2*(h_tv/count_h+w_tv/count_w)/batch_size

    @staticmethod
    def _tensor_size(t):
        return t.size()[1]*t.size()[2]*t.size()[3]


class ConvBlock(nn.Module):

    def __init__(self):
        super(ConvBlock, self).__init__()
        self.conv1 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
        self.instance_norm1 = nn.InstanceNorm2d(64, affine=True)
        self.instance_norm2 = nn.InstanceNorm2d(64, affine=True)
        self.relu = nn.ReLU(inplace=True)

    def forward(self, x):
        y = self.relu(self.instance_norm1(self.conv1(x)))
        y = self.instance_norm2(self.conv2(y)) + x
        return y


class Enhance_Generator(nn.Module):

    def __init__(self):
        super(Enhance_Generator, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.blocks = nn.Sequential(
            ConvBlock(),
            ConvBlock(),
            ConvBlock(),
            ConvBlock(),
            ConvBlock(),
            ConvBlock(),
            ConvBlock(),
            ConvBlock(),
            ConvBlock(),
        )
        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
        self.instance_norm = nn.InstanceNorm2d(64, affine=True)
        self.conv3 = nn.Conv2d(64, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 3, 1, padding=0)

    def forward(self, x):

        x = F.relu(self.conv1(x))
        temp = x
        x = self.blocks(x)
        x = self.instance_norm(self.conv2(x)) + temp
        x = F.relu(self.conv3(x))
        x = F.tanh(self.conv4(x))

        return x


class Enhance_DiscriminatorC(nn.Module):

    def __init__(self, input_ch):
        super(Enhance_DiscriminatorC, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(input_ch, 48, 11, stride=4, padding=5),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.Conv2d(48, 128, 5, stride=2, padding=2),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.InstanceNorm2d(128, affine=True),
            nn.Conv2d(128, 192, 3, stride=1, padding=1),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.InstanceNorm2d(192, affine=True),
            nn.Conv2d(192, 192, 3, stride=1, padding=1),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.InstanceNorm2d(192, affine=True),
            nn.Conv2d(192, 128, 3, stride=2, padding=1),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.InstanceNorm2d(128, affine=True),
        )

        self.fc = nn.Linear(128*9*9, 1024)
        self.out = nn.Linear(1024, 2)

    def forward(self, x):
        x = self.conv_layers(x)
        #print(x.shape)
        x = x.view(-1, 128*9*9)
        #print(x.shape)
        x = F.leaky_relu(self.fc(x), negative_slope=0.2)
        x = F.softmax(self.out(x))
        return x


class Enhance_DiscriminatorT(nn.Module):

    def __init__(self, input_ch):
        super(Enhance_DiscriminatorT, self).__init__()
        self.conv_layers = nn.Sequential(
            nn.Conv2d(input_ch, 48, 11, stride=4, padding=5),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.Conv2d(48, 128, 5, stride=2, padding=2),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.InstanceNorm2d(128, affine=True),
            nn.Conv2d(128, 192, 3, stride=1, padding=1),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.InstanceNorm2d(192, affine=True),
            nn.Conv2d(192, 192, 3, stride=1, padding=1),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.InstanceNorm2d(192, affine=True),
            nn.Conv2d(192, 128, 3, stride=2, padding=1),
            nn.LeakyReLU(negative_slope=0.2, inplace=True),
            nn.InstanceNorm2d(128, affine=True),
        )

        self.fc = nn.Linear(128*8*8, 1024)
        self.out = nn.Linear(1024, 2)

    def forward(self, x):
        x = self.conv_layers(x)
        #print(x.shape)
        x = x.view(-1, 128*8*8)
        #print(x.shape)
        x = F.leaky_relu(self.fc(x), negative_slope=0.2)
        x = F.softmax(self.out(x))
        return x


class GaussianBlur(nn.Module):
    def __init__(self):
        super(GaussianBlur, self).__init__()
        kernel = [[0.03797616, 0.044863533, 0.03797616],
                  [0.044863533, 0.053, 0.044863533],
                  [0.03797616, 0.044863533, 0.03797616]]
        kernel = torch.FloatTensor(kernel).unsqueeze(0).unsqueeze(0)
        # self.weight = nn.Parameter(data=kernel, requires_grad=False)
        self.weight = kernel.cuda()

    def forward(self, x):
        x1 = x[:, 0]
        x2 = x[:, 1]
        x3 = x[:, 2]
        x1 = F.conv2d(x1.unsqueeze(1), self.weight, padding=2)
        x2 = F.conv2d(x2.unsqueeze(1), self.weight, padding=2)
        x3 = F.conv2d(x3.unsqueeze(1), self.weight, padding=2)
        x = torch.cat([x1, x2, x3], dim=1)
        return x


class GrayLayer(nn.Module):

    def __init__(self):
        super(GrayLayer, self).__init__()

    def forward(self, x):
        result = 0.299 * x[:, 0] + 0.587 * x[:, 1] + 0.114 * x[:, 2]
        # print(result.unsqueeze(1).shape)
        return result.unsqueeze(1)


class VGG(nn.Module):

    def __init__(self):
        super(VGG, self).__init__()
        self.model = vgg19_bn(True).features.cuda()
        self.mean = torch.Tensor([123.68,  116.779,  103.939]).cuda().view(1,3,1,1)
        for param in self.model.parameters():
            param.requires_grad = False

    def forward(self, x):
        x = x*255 - self.mean
        #print(type(x), x.shape)
        x = self.model(x)
        return x
<file_sep>/README.md
# SR_2020
A python vision for our lastest face super-resolution method.

#### TODO:The code will be updated until it is available and normalized.

Hoping the COVID-19 will end soon.
