text: # PyTorch lib
  import torch
  import torch.nn as nn
  from torch.autograd import Variable
  import torch.utils.data as Data
  import torch.nn.functional as F
  import torchvision
  # Tools lib
  import numpy as np
  import cv2
  import random
  import time
  import os
  
  # Set iteration time
  ITERATION = 4
  
  
  # Model
  class Att_GAN_Generator(nn.Module):
      def __init__(self):
          super(Att_GAN_Generator, self).__init__()
          self.det_conv0 = nn.Sequential(
              nn.Conv2d(4, 32, 3, 1, 1),
              nn.ReLU()
          )
          self.det_conv1 = nn.Sequential(
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU(),
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU()
          )
          self.det_conv2 = nn.Sequential(
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU(),
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU()
          )
          self.det_conv3 = nn.Sequential(
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU(),
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU()
          )
          self.det_conv4 = nn.Sequential(
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU(),
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU()
          )
          self.det_conv5 = nn.Sequential(
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU(),
              nn.Conv2d(32, 32, 3, 1, 1),
              nn.ReLU()
          )
          self.conv_i = nn.Sequential(
              nn.Conv2d(32 + 32, 32, 3, 1, 1),
              nn.Sigmoid()
          )
          self.conv_f = nn.Sequential(
              nn.Conv2d(32 + 32, 32, 3, 1, 1),
              nn.Sigmoid()
          )
          self.conv_g = nn.Sequential(
              nn.Conv2d(32 + 32, 32, 3, 1, 1),
              nn.Tanh()
          )
          self.conv_o = nn.Sequential(
              nn.Conv2d(32 + 32, 32, 3, 1, 1),
              nn.Sigmoid()
          )
          self.det_conv_mask = nn.Sequential(
              nn.Conv2d(32, 1, 3, 1, 1),
          )
          self.conv1 = nn.Sequential(
              nn.Conv2d(4, 64, 5, 1, 2),
              nn.ReLU()
          )
          self.conv2 = nn.Sequential(
              nn.Conv2d(64, 128, 3, 2, 1),
              nn.ReLU(),
              nn.Conv2d(128, 128, 3, 1, 1),
              nn.ReLU()
          )
          self.conv4 = nn.Sequential(
              nn.Conv2d(128, 256, 3, 2, 1),
              nn.ReLU(),
              nn.Conv2d(256, 256, 3, 1, 1),
              nn.ReLU(),
              nn.Conv2d(256, 256, 3, 1, 1),
              nn.ReLU()
          )
          self.diconv1 = nn.Sequential(
              nn.Conv2d(256, 256, 3, 1, 2, dilation=2),
              nn.ReLU(),
              nn.Conv2d(256, 256, 3, 1, 4, dilation=4),
              nn.ReLU(),
              nn.Conv2d(256, 256, 3, 1, 8, dilation=8),
              nn.ReLU(),
              nn.Conv2d(256, 256, 3, 1, 16, dilation=16),
              nn.ReLU()
          )
          self.conv7 = nn.Sequential(
              nn.Conv2d(256, 256, 3, 1, 1),
              nn.ReLU(),
              nn.Conv2d(256, 256, 3, 1, 1),
              nn.ReLU()
          )
          self.deconv1 = nn.Sequential(
              nn.ConvTranspose2d(256, 128, 4, 2, 1),
              nn.ReflectionPad2d((1, 0, 1, 0)),
              nn.AvgPool2d(2, stride=1),
              nn.ReLU()
          )
          self.conv9 = nn.Sequential(
              nn.Conv2d(128, 128, 3, 1, 1),
              nn.ReLU()
          )
          self.deconv2 = nn.Sequential(
              nn.ConvTranspose2d(128, 64, 4, 2, 1),
              nn.ReflectionPad2d((1, 0, 1, 0)),
              nn.AvgPool2d(2, stride=1),
              nn.ReLU()
          )
          self.conv10 = nn.Sequential(
              nn.Conv2d(64, 32, 3, 1, 1),
              nn.ReLU()
          )
          self.outframe1 = nn.Sequential(
              nn.Conv2d(256, 3, 3, 1, 1),
              nn.ReLU()
          )
          self.outframe2 = nn.Sequential(
              nn.Conv2d(128, 3, 3, 1, 1),
              nn.ReLU()
          )
          self.output = nn.Sequential(
              nn.Conv2d(32, 3, 3, 1, 1)
          )
  
      def forward(self, input):
          batch_size, row, col = input.size(0), input.size(2), input.size(3)
          mask = Variable(torch.ones(batch_size, 1, row, col)).cuda() / 2.
          h = Variable(torch.zeros(batch_size, 32, row, col)).cuda()
          c = Variable(torch.zeros(batch_size, 32, row, col)).cuda()
          mask_list = []
          for i in range(ITERATION):
              x = torch.cat((input, mask), 1)
              x = self.det_conv0(x)
              resx = x
              x = F.relu(self.det_conv1(x) + resx)
              resx = x
              x = F.relu(self.det_conv2(x) + resx)
              resx = x
              x = F.relu(self.det_conv3(x) + resx)
              resx = x
              x = F.relu(self.det_conv4(x) + resx)
              resx = x
              x = F.relu(self.det_conv5(x) + resx)
              x = torch.cat((x, h), 1)
              i = self.conv_i(x)
              f = self.conv_f(x)
              g = self.conv_g(x)
              o = self.conv_o(x)
              c = f * c + i * g
              h = o * F.tanh(c)
              mask = self.det_conv_mask(h)
              mask_list.append(mask)
          x = torch.cat((input, mask), 1)
          x = self.conv1(x)
          res1 = x
          x = self.conv2(x)
          res2 = x
          x = self.conv4(x)
          x = self.diconv1(x)
          x = self.conv7(x)
          frame1 = self.outframe1(x)
          x = self.deconv1(x)
          x = x + res2
          x = self.conv9(x)
          frame2 = self.outframe2(x)
          x = self.deconv2(x)
          x = x + res1
          x = self.conv10(x)
          x = self.output(x)
          return mask_list, frame1, frame2, x
  
  
  class Att_GAN_Discriminator(nn.Module):
      def __init__(self):
          super(Att_GAN_Discriminator, self).__init__()
          self.conv1 = nn.Sequential(
              nn.Conv2d(3, 8, 5, 1, 2),
              nn.ReLU(),
              nn.Conv2d(8, 16, 5, 1, 2),
              nn.ReLU(),
              nn.Conv2d(16, 64, 5, 1, 2),
              nn.ReLU(),
              nn.Conv2d(64, 128, 5, 1, 2),
              nn.ReLU(),
              nn.Conv2d(128, 128, 5, 1, 2),
              nn.ReLU(),
              nn.Conv2d(128, 128, 5, 1, 2),
              nn.ReLU()
          )
          self.conv_mask = nn.Sequential(
              nn.Conv2d(128, 1, 5, 1, 2)
          )
          self.conv7 = nn.Sequential(
              nn.Conv2d(128, 64, 5, 4, 1),
              nn.ReLU()
          )
          self.conv8 = nn.Sequential(
              nn.Conv2d(64, 32, 5, 4, 1),
              nn.ReLU()
          )
          self.fc = nn.Sequential(
              nn.Linear(32 * 14 * 14, 1024),
              nn.Linear(1024, 1),
              nn.Sigmoid()
          )
  
      def forward(self, x):
          x = self.conv1(x)
          mask = self.conv_mask(x)
          x = self.conv7(x * mask)
          x = self.conv7(x)
          x = self.conv8(x)
          x = x.view(x.size(0), -1)
          return mask, self.fc(x)
  
  
  def trainable(net, trainable):
      for para in net.parameters():
          para.requires_grad = trainable
  
  
  # Initialize VGG16 with pretrained weight on ImageNet
  def vgg_init():
      vgg_model = torchvision.models.vgg16(pretrained=True).cuda()
      trainable(vgg_model, False)
      return vgg_model
  
  
  # Extract features from internal layers for perceptual loss
  class vgg(nn.Module):
      def __init__(self, vgg_model):
          super(vgg, self).__init__()
          self.vgg_layers = vgg_model.features
          self.layer_name_mapping = {
              '1': "relu1_1",
              '3': "relu1_2",
              '6': "relu2_1",
              '8': "relu2_2"
          }
  
      def forward(self, x):
          output = []
          for name, module in self.vgg_layers._modules.items():
              x = module(x)
              if name in self.layer_name_mapping:
                  output.append(x)
          return output
  