<repo_name>Cl2rk3Gr1ff1n/Synthesis-of-Frontal-Face-Disentangled-Poses-and-Illuminations<file_sep>/siamese_model_2nd.py
import torch.nn as nn
import torch
import torch.nn.functional as f
class Siamese(nn.Module):
    def __init__(self):
        super(Siamese,self).__init__()
        # 3x96x96
        self.conv11 = nn.Sequential(nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(64),
                                    nn.ELU(True))
        # 3x96x96 conv1r has RES
        self.conv1r = nn.Sequential(nn.Conv2d(3, 64, kernel_size=1, stride=1, padding=0))
        # 64x96x96 
        self.conv12 = nn.Sequential(nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(64))
        # sum with residual, then pass the ELU block
        self.elu12 = nn.ELU(True)
        #-------------------------------------------1st layer
        # 64x96x96 conv21 has RES
        self.conv21 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),
                                    nn.BatchNorm2d(128),
                                    nn.ELU(True))
        # 128x48x48
        self.conv22 = nn.Sequential(nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(64),
                                    nn.ELU(True))
        # 64x48x48
        self.conv23 = nn.Sequential(nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(128))
        # sum with residual, then pass to ELU block
        self.elu23 = nn.ELU(True)
        #-------------------------------------------2nd layer
        # 128x48x48 conv31 has RES
        self.conv31 = nn.Sequential(nn.Conv2d(128, 192, kernel_size=3, stride=2, padding=1),
                                    nn.BatchNorm2d(192),
                                    nn.ELU(True))
        # 192x24x24
        self.conv32 = nn.Sequential(nn.Conv2d(192, 96, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(96),
                                    nn.ELU(True))
        # 96x24x24
        self.conv33 = nn.Sequential(nn.Conv2d(96, 192, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(192))
        # sum with residual, then pass to ELU block
        self.elu33 = nn.ELU(True)
        #-------------------------------------------3rd layer
        #192x24x24 conv41 has RES
        self.conv41 = nn.Sequential(nn.Conv2d(192, 256, kernel_size=3, stride=2, padding=1),
                                    nn.BatchNorm2d(256),
                                    nn.ELU(True))
        #256x12x12
        self.conv42 = nn.Sequential(nn.Conv2d(256, 128, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(128),
                                    nn.ELU(True))
        #128x12x12
        self.conv43 = nn.Sequential(nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(256))
        # sum with residual, then pass to ELU block
        self.elu43 = nn.ELU(True)
        #-------------------------------------------4th layer
        #256x12x12
        self.conv51 = nn.Sequential(nn.Conv2d(256, 320, kernel_size=3, stride=2, padding=1),
                                    nn.BatchNorm2d(320),
                                    nn.ELU(True))
        #320x6x6
        self.conv52 = nn.Sequential(nn.Conv2d(320, 160, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(160),
                                    nn.ELU(True))
        #160x6x6
        self.conv53 = nn.Sequential(nn.Conv2d(160, 320, kernel_size=3, stride=1, padding=1),
                                    nn.BatchNorm2d(320))
        # sum with residual, then pass to ELU block
        self.elu53 = nn.ELU(True)
        #output size 320x6x6
        #-------------------------------------------5th layer
        #-------------------------------------------decomposition
        #320x6x6
        # view->240 + 80
        #
        #240x6x6 Image_1 and Image_2 compute contrastive loss

        # 80 pas through conv to be 27
        # 80x6x6
        self.convfc = nn.Conv2d(80, 29, kernel_size=6, stride=1, padding=0)
        #27x1x1
        #27 slice into 20x1x1 + 7x1x1
        # self.fc_pose = nn.Linear(27 , 20)
        # self.fc_light = nn.Linear(27, 7)
        #compute contractiveLoss
        #compute Softmax of pose and light
        #predict Label of pose and light
        #-------------------------------------------decoder-------------------------------------------------------------------
        #240x6x6
        self.dconv52 = nn.Sequential(nn.ConvTranspose2d(240, 160, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(160),
                                     nn.ELU(True))
        #160x6x6
        self.dconv51 = nn.Sequential(nn.ConvTranspose2d(160, 256, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(256),
                                     nn.ELU(True))
        #256x6x6 bilinear interpolation upsamping
        #self.upsampling43 = nn.UpsamplingBilinear2d(scale_factor=2)
        #256x12x12
        self.dconv43 = nn.Sequential(nn.ConvTranspose2d(256 , 256, kernel_size = 3, stride=2, padding=1, output_padding=1),
                                     nn.BatchNorm2d(256),
                                     nn.ELU(True))
        #256x12x12
        self.dconv42 = nn.Sequential(nn.ConvTranspose2d(256, 128, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(128),
                                     nn.ELU(True))
        #128x12x12
        self.dconv41 = nn.Sequential(nn.ConvTranspose2d(128, 192, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(192),
                                     nn.ELU(True))
        #192x12x12 bilinear interpolation upsamping
        #self.upsampling33 = nn.UpsamplingBilinear2d(scale_factor=2)
        #192x24x24
        self.dconv33 = nn.Sequential(nn.ConvTranspose2d(192 , 192, kernel_size = 3, stride=2, padding=1, output_padding=1),
                                     nn.BatchNorm2d(192),
                                     nn.ELU(True))
        #192x24x24
        self.dconv32 = nn.Sequential(nn.ConvTranspose2d(192, 96, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(96),
                                     nn.ELU(True))
        #96x24x24
        self.dconv31 = nn.Sequential(nn.ConvTranspose2d(96, 128, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(128),
                                     nn.ELU(True))
        #128x24x24 bilinear interpolation upsamping
        #self.upsampling23 = nn.UpsamplingBilinear2d(scale_factor=2)
        #128x48x48
        self.dconv23 = nn.Sequential(nn.ConvTranspose2d(128 , 128, kernel_size = 3, stride=2, padding=1, output_padding=1),
                                     nn.BatchNorm2d(128),
                                     nn.ELU(True))
        #128x48x48
        self.dconv22 = nn.Sequential(nn.ConvTranspose2d(128, 64, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(64),
                                     nn.ELU(True))
        #64x48x48
        self.dconv21 = nn.Sequential(nn.ConvTranspose2d(64, 64, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(64),
                                     nn.ELU(True))
        #64x48x48 bilinear interpolation upsamping
        #self.upsampling13 = nn.UpsamplingBilinear2d(scale_factor=2)
        #64x96x96
        self.dconv12 = nn.Sequential(nn.ConvTranspose2d(64 , 64, kernel_size = 3, stride=2, padding=1, output_padding=1),
                                     nn.BatchNorm2d(64),
                                     nn.ELU(True))
        #64x96x96
        self.dconv11 = nn.Sequential(nn.ConvTranspose2d(64, 32, kernel_size=3, stride=1, padding=1),
                                     nn.BatchNorm2d(32),
                                     nn.ELU(True))
        #32x96x96
        self.output = nn.Sequential(nn.ConvTranspose2d(32, 3, kernel_size=3, stride=1, padding=1),
                                    nn.Tanh())
        #3x96x96
        #self.tanh = nn.tanh()
        #compute the L1 loss(Lrecon) between GT and reconstruction 
    def forward(self, x, x_2):
        #-----layer 1------------------------------
        out_res_1 = self.conv1r(x)
        out = self.conv11(x)
        out = self.conv12(out)
        out = out + out_res_1
        out = self.elu12(out)

        out_res_1_2 = self.conv1r(x_2)
        out_2 = self.conv11(x_2)
        out_2 = self.conv12(out_2)
        out_2 = out_2 + out_res_1_2
        out_2 = self.elu12(out_2)
        #-----layer 2------------------------------
        out_res_2 = self.conv21(out)
        out = self.conv22(out_res_2)
        out = self.conv23(out)
        out = out + out_res_2
        out = self.elu23(out)

        out_res_2_2 = self.conv21(out_2)
        out_2 = self.conv22(out_res_2_2)
        out_2 = self.conv23(out_2)
        out_2 = out_2 + out_res_2_2
        out_2 = self.elu23(out_2)
        #-----layer 3-------------------------------
        out_res_3 = self.conv31(out)
        out = self.conv32(out_res_3)
        out = self.conv33(out)
        out = out + out_res_3
        out = self.elu33(out)

        out_res_3_2 = self.conv31(out_2)
        out_2 = self.conv32(out_res_3_2)
        out_2 = self.conv33(out_2)
        out_2 = out_2 + out_res_3_2
        out_2 = self.elu33(out_2)
        #-----layer 4--------------------------------
        out_res_4 = self.conv41(out)
        out = self.conv42(out_res_4)
        out = self.conv43(out)
        out = out + out_res_4
        out = self.elu43(out)

        out_res_4_2 = self.conv41(out_2)
        out_2 = self.conv42(out_res_4_2)
        out_2 = self.conv43(out_2)
        out_2 = out_2 + out_res_4_2
        out_2 = self.elu43(out_2)
        #-----layer 5--------------------------------
        out_res_5 = self.conv51(out)
        out = self.conv52(out_res_5)
        out = self.conv53(out)
        out = out + out_res_5
        out = self.elu53(out)

        out_res_5_2 = self.conv51(out_2)
        out_2 = self.conv52(out_res_5_2)
        out_2 = self.conv53(out_2)
        out_2 = out_2 + out_res_5_2
        out_2 = self.elu53(out_2)
        #----------------------------------------decomposition----------------------------------------------
        out_240 = out.narrow(1,0,240) #second dimension[channels]slices from 0 to 239
        out_80 = out.narrow(1,240,80) #second dimension[channels]slices from 240 to 319
        out_240_2 = out_2.narrow(1, 0, 240)
        out_80_2 = out_2.narrow(1, 240, 80)

        out_29 = self.convfc(out_80)
        out_29_2 = self.convfc(out_80_2)
        #29x1x1
        out_pose = out_29.narrow(1,0,9) #second dimension[channels]slices from 0 to 6 for pose

        #9x1x1
        out_pose_2 = out_29_2.narrow(1, 0, 9)


        out_light = out_29.narrow(1,9,20) #second dimension[channels]slices from 7 to 26 for light
        #20x1x1

        out_light_2 = out_29_2.narrow(1, 9 , 20)


        #-----------------------------------------decoder---------------------------------------------------
        out = self.dconv52(out_240)
        out = self.dconv51(out)
        
        # out = f.upsample(out, scale_factor = 2, mode='bilinear')
        out = self.dconv43(out)
        out = self.dconv42(out)
        out = self.dconv41(out)
        
        # out = f.upsample(out, scale_factor = 2, mode='bilinear')
        out = self.dconv33(out)
        out = self.dconv32(out)
        out = self.dconv31(out)
        
        # out = f.upsample(out, scale_factor = 2, mode='bilinear')
        out = self.dconv23(out)
        out = self.dconv22(out)
        out = self.dconv21(out)

        # out = f.upsample(out, scale_factor = 2, mode='bilinear')
        out = self.dconv12(out)
        out = self.dconv11(out)
        out = self.output(out)
        # out = out.tanh()

        out_pose = out_pose.contiguous()
        out_pose_2 = out_pose_2.contiguous()
        out_light = out_light.contiguous()
        out_light_2 = out_light_2.contiguous()

        out_pose = out_pose.view(out_pose.size(0),-1)
        out_pose_2 = out_pose_2.view(out_pose_2.size(0),-1)
        out_light = out_light.view(out_light.size(0),-1)
        out_light_2 = out_light_2.view(out_light_2.size(0),-1)

        out_240 = out_240.contiguous()
        out_240 = out_240.view(64, -1)
        # 64x8640
        out_240_2 = out_240_2.contiguous()
        out_240_2 = out_240_2.view(64, -1)
        #--------------------------------------------LOSS
        # contractive loss between gt and image_1's identity
        return out_pose, out_pose_2, out_light, out_light_2, out_240, out_240_2, out
















<file_sep>/mask/multi_mask.py
from PIL import Image
from pylab import *
import numpy as np

im = Image.open('249_01_01_051_00.png')
im = im.resize((96, 96), Image.BILINEAR)
# width, height = im.size
# print width,height
# 96x96
img = array(im)
imshow(img)

coordinate = ginput(300)
mask = np.load('mask_19.npy')

for i in coordinate:
    y = int(i[0])
    x = int(i[1])
    mask[x, y, 0] = 5
    mask[x, y, 1] = 5
    mask[x, y, 2] = 5
    # use print img[x, y, 0] & print img[1,60:90,:] to check
    # the pixel location and coordinate is inversed
    # coordinate(82,1) right top but pixel[1,82,:] means right top
np.save("mask_20", mask)
print mask
show()


<file_sep>/README.md
# Siamese-Network-for-Frontal-Face-Synthesis-disentangle-pose-and-light-interference-

All the input data are from MultiPIE dataset, 9 different identity faces each time with different poses and illumination.
![input](https://github.com/danny95333/Synthesis-of-Frontal-Face-Disentangled-Poses-and-Illuminations/blob/master/input_samples_iteration_200.png)

## First version: complete basic functions, the frontal face image has clear identity, but blurry boundary and glasses still can't be 100% synthesis.
>siamese_out_largeP&LclassW: these are the output with relative large class_P and class_L LOSS.
![first_L](https://github.com/danny95333/Synthesis-of-Frontal-Face-Disentangled-Poses-and-Illuminations/blob/master/siamese_out_largeP%26LclassW/fake_samples_iteration_70000.png)
>siamese_out_sameW: these are the output with same weigths to every LOSSes.
![first_Same](https://github.com/danny95333/Synthesis-of-Frontal-Face-Disentangled-Poses-and-Illuminations/blob/master/siamese_out_sameW/fake_samples_iteration_70000.png)
## 2nd update at 30th June
* load from June-20 pretrained model and get rid of other losses but L1 loss
>2nd_step_only_L1: these are the 2nd update's output with 'Step_decay_learning_rate'
![second_l1](https://github.com/danny95333/Synthesis-of-Frontal-Face-Disentangled-Poses-and-Illuminations/blob/master/2nd_step_only_L1/fake_samples_iteration_40000.png)
## 3rd update: Mask was sample from 20 frontal face image manually, the model and weights was based on the 1st version.
* Add Mask at frontal face's boundary to make the netS be more sensitive to the boundary region
* Freeze the pre-trained Siamese net's encoder part, which aimed to capture the input images' feature maps. And train the decoder part of the network.
>3rd_step_L1_Mask_FreezEnc: These are the 3rd update's output with 'Step_decay_learning_rate', still using L1 as general supervise loss.
![third_l1](https://github.com/danny95333/Synthesis-of-Frontal-Face-Disentangled-Poses-and-Illuminations/blob/master/3rd_step_L1_Mask_FreezEnc/fake_samples_iteration_30000.png)
>3rd_step_L2_Mask_FreezEnc: These are the 3rd update's output with 'Step_decay_learning_rate', still using L2 as general supervise loss.
Problem: There are still many artifacts on the synthesized images, not 'real' enough 
![third_l2](https://github.com/danny95333/Synthesis-of-Frontal-Face-Disentangled-Poses-and-Illuminations/blob/master/3rd_step_L2_Mask_FreezEnc/fake_samples_iteration_30000.png)



<file_sep>/siamese_train_M_3.py
#from __future__ import print_function
import cv2
import argparse
import os
import random
import torch
import torch.nn as nn
import torch.backends.cudnn as cudnn
import torch.optim as optim
import torchnet.meter as meter
import torchvision.datasets as dset
import torchvision.transforms as transforms
import torchvision.utils as vutils
from torch.autograd import Variable
from dataset import multiPIE
from siamese_model_2nd import Siamese
from contrastive import ContrastiveLoss
import numpy as np
# import cv2
#from pycrayon import CrayonClient
 
#for plotting loss
import matplotlib
matplotlib.use('agg')
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
import time,math
from logger import Logger
# from models_Parsing import ParseNet
saveFile = open('/home/shumao/wyw_files/siamese_output_M_3/record.txt', 'w')
saveFile.write("niter:" + str(50000) + "\n")
saveFile.write("---lr:" + str(0.0001) + "\n")
saveFile.write("beta1:" + str(0.7) + "\n")
saveFile.write("W:-1-x-x-x-x-x-" + "\n")
saveFile.write("L1 Loss" + "\n")
saveFile.write("after load model from: train-3-28000pth")
logger = Logger('./log_1');

parser = argparse.ArgumentParser()
parser.add_argument('--batchSize', type=int, default=64, help='input batch size')
parser.add_argument('--loadSize', type=int, default=100, help='the height / width of the input image to network')
parser.add_argument('--fineSize', type=int, default=96, help='the height / width of the input image to network')
parser.add_argument('--id_num', type=int, default=200, help='Total training identity.')
parser.add_argument('--pose_num', type=int, default=9, help='Total training pose.')
parser.add_argument('--light_num', type=int, default=20, help='Total training light.')
parser.add_argument('--niter', type=int, default=50000, help='number of iterations to train for')
parser.add_argument('--lr', type=float, default=0.0001, help='learning rate, default=0.0002')
parser.add_argument('--beta1', type=float, default=0.7, help='beta1 for adam. default=0.7')
parser.add_argument('--cuda', action='store_true', help='enables cuda')
parser.add_argument('--outf', default='/home/shumao/wyw_files/siamese_output_M_3', help='folder to output images and model checkpoints')
parser.add_argument('--manualSeed', type=int, help='manual seed')
parser.add_argument('--dataPath', default='/home/shumao/dr-gan/Data_new_realigned2/setting2/train/', help='which dataset to train on')
parser.add_argument('--modelPath', default='/home/shumao/wyw_files/siamese_output_3/netS_28000.pth', help='which dataset to train on')
parser.add_argument('--save_step', type=int, default=400, help='save weights every 400 iterations ')
parser.add_argument('--labelPath', default='/home/shumao/dr-gan/Data_new_realigned2/setting2/Facedata/', help='which dataset to train on')


opt = parser.parse_args()
print(opt) # print every parser arguments
# print(opt.niter)


try:
    os.makedirs(opt.outf)
except OSError:
    pass

w_r = 1
# w_cL = 0.02
# w_cP = 0.02
# w_cI = 0.02
# w_P = 0.02
# w_L = 0.02

if opt.manualSeed is None:
    opt.manualSeed = random.randint(1, 10000)
print("Random Seed: ", opt.manualSeed)
random.seed(opt.manualSeed)
torch.manual_seed(opt.manualSeed)
if opt.cuda:
    torch.cuda.manual_seed_all(opt.manualSeed)

cudnn.benchmark = True
#---------------------Load Mask-------------------
mask = np.load('mask_20.npy')
mask = mask.astype(np.float32)
M = torch.from_numpy(mask.transpose((2, 0, 1)))
FinalMask = M.expand(opt.batchSize,3,96,96)
# print m.size()
# 3x96x96


#---------------------Load DATA-------------------------
dataset_1 = multiPIE(opt.dataPath,opt.loadSize,opt.fineSize,labelPath = opt.labelPath)
# dataset_2 = multiPIE(opt.dataPath,opt.loadSize,opt.fineSize,opt.labelPath)
dataset_test = multiPIE('/home/shumao/dr-gan/comparison/',opt.loadSize,opt.fineSize,labelPath = opt.labelPath)
loader_train_1 = torch.utils.data.DataLoader(dataset=dataset_1,
                                      batch_size = opt.batchSize,
                                      shuffle=True,
                                      num_workers=4,
                                      drop_last = True)
# loader_train_2 = torch.utils.data.Dataloader(dataset=dataset_1,
#                                       batch_size = opt.batchSize,
#                                       shuffle=True,
#                                       num_workers=4)


loader_test = torch.utils.data.DataLoader(dataset=dataset_test,
                                          batch_size = 9,
                                          shuffle=False,
                                          num_workers=4)
data_train_1 = iter(loader_train_1)
# data_train_2 = iter(loader_train_2)
data_test = iter(loader_test)


#----------------------Parameters-----------------------
num_pose = opt.pose_num
num_light = opt.light_num
num_iden = opt.id_num


def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') !=-1:
        m.weight.data.normal_(0.0, 0.02)
    elif classname.find('BatchNorm') !=-1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)




netS = Siamese()
netS.load_state_dict(torch.load(opt.modelPath))

#-----------------params freeze-----------------
for param in netS.conv11.parameters():
    param.requires_grad = False
for param in netS.conv1r.parameters():
    param.requires_grad = False
for param in netS.conv12.parameters():
    param.requires_grad = False
for param in netS.conv21.parameters():
    param.requires_grad = False
for param in netS.conv22.parameters():
    param.requires_grad = False
for param in netS.conv23.parameters():
    param.requires_grad = False
for param in netS.conv31.parameters():
    param.requires_grad = False
for param in netS.conv32.parameters():
    param.requires_grad = False
for param in netS.conv33.parameters():
    param.requires_grad = False
for param in netS.conv41.parameters():
    param.requires_grad = False
for param in netS.conv42.parameters():
    param.requires_grad = False
for param in netS.conv43.parameters():
    param.requires_grad = False
for param in netS.conv51.parameters():
    param.requires_grad = False
for param in netS.conv52.parameters():
    param.requires_grad = False
for param in netS.conv53.parameters():
    param.requires_grad = False
for param in netS.convfc.parameters():
    param.requires_grad = False


#-----------------params freeze-----------------
if(opt.cuda):
    netS.cuda()

#-------------------Loss & Optimization

optimizerS = torch.optim.Adam(filter(lambda p: p.requires_grad, netS.parameters()),lr=opt.lr, betas=(opt.beta1, 0.999))

poss_contrastive_loss = ContrastiveLoss() # load from the begining
light_contrastive_loss = ContrastiveLoss()
identity_contrastive_loss = ContrastiveLoss()
reconstructe_loss = nn.L1Loss()
pose_class_loss = nn.CrossEntropyLoss()
light_class_loss = nn.CrossEntropyLoss()

#------------------ Global Variables------------------
input_pose_1 = torch.LongTensor(opt.batchSize)
input_light_1 = torch.LongTensor(opt.batchSize)
# input_pose_2 = torch.LongTensor(opt.batchSize)
# input_light_2 = torch.LongTensor(opt.batchSize)

inputImg_1 = torch.FloatTensor(opt.batchSize, 3, opt.fineSize, opt.fineSize)
inputImg_2 = torch.FloatTensor(opt.batchSize, 3, opt.fineSize, opt.fineSize)
GT = torch.FloatTensor(opt.batchSize, 3,opt.fineSize, opt.fineSize)
same_pose = torch.FloatTensor(opt.batchSize)
same_iden = torch.FloatTensor(opt.batchSize)
same_light = torch.FloatTensor(opt.batchSize)

# w_1 = torch.FloatTensor(1)
# w_2 = torch.FloatTensor(20)
# w_3 = torch.FloatTensor(10)
# w_4 = torch.FloatTensor(10)
# w_5 = torch.FloatTensor(10)
# w_6 = torch.FloatTensor(20)
# output_pose_1_label = torch.LongTensor(opt.batchSize)
# output_pose_2_label = torch.LongTensor(opt.batchSize)
# output_light_1_label = torch.LongTensor(opt.batchSize)
# output_light_2_label = torch.LongTensor(opt.batchSize)

input_pose_1 = Variable(input_pose_1)
# input_pose_2 = Variable(input_pose_2)
input_light_1 = Variable(input_light_1)
# input_light_2 = Variable(input_light_2)

inputImg_1 = Variable(inputImg_1)
inputImg_2 = Variable(inputImg_2)
GT = Variable(GT)
same_pose = Variable(same_pose)
same_iden = Variable(same_iden)
same_light = Variable(same_light)

FinalMask = Variable(FinalMask)

# w_1 = Variable(w_1, requires_grad = False)
# w_2 = Variable(w_2, requires_grad = False)
# w_3 = Variable(w_3, requires_grad = False)
# w_4 = Variable(w_4, requires_grad = False)
# w_5 = Variable(w_5, requires_grad = False)
# w_6 = Variable(w_6, requires_grad = False)


pose_mtr = meter.ConfusionMeter(k=opt.pose_num)
light_mtr = meter.ConfusionMeter(k=opt.light_num)


if(opt.cuda):

    input_pose_1 = input_pose_1.cuda()
    # input_pose_2 = input_pose_2.cuda()
    input_light_1 = input_light_1.cuda()
    # input_light_2 = input_light_2.cuda()
    inputImg_1 = inputImg_1.cuda()
    inputImg_2 = inputImg_2.cuda()
    GT = GT.cuda()
    same_pose = same_pose.cuda()
    same_light = same_light.cuda()
    same_iden = same_iden.cuda()

    FinalMask = FinalMask.cuda()

    # w_1 = w_1.cuda()
    # w_2 = w_1.cuda()
    # w_3 = w_1.cuda()
    # w_4 = w_1.cuda()
    # w_5 = w_1.cuda()
    # w_6 = w_1.cuda()
    # poss_contrastive_loss.cuda()
    # light_contrastive_loss.cuda()
    # identity_contrastive_loss.cuda()
    pose_class_loss.cuda()
    light_class_loss.cuda()
    reconstructe_loss.cuda()


#------------------test---------

# k = 0 # for meter

err_total = 0
err_recon = 0
err_contraL = 0
err_contraP = 0
err_contraI = 0
err_classP = 0
err_classL = 0

def test(iteration, data_test, loader_test):
    try:
        images_1,po_1,li_1,GT_1,by_image,same_po,same_li,same_id = data_test.next()
    except StopIteration:
        data_test = iter(loader_test)
        images_1,po_1,li_1,GT_1,by_image,same_po,same_li,same_id = data_test.next()

    GT.data.resize_(GT_1.size()).copy_(GT_1)
    inputImg_1.data.resize_(images_1.size()).copy_(images_1)
    inputImg_2.data.resize_(by_image.size()).copy_(by_image)
    input_pose_1.data.resize_(po_1.size()).copy_(po_1)
    input_light_1.data.resize_(li_1.size()).copy_(li_1)


    output_pose_1, output_pose_2, output_light_1, output_light_2, out_f_1, out_f_2, out = netS(inputImg_1, inputImg_2)
    vutils.save_image(out.data,
        '%s/fake_samples_iteration_%03d.png' % (opt.outf, iteration), normalize=True)
    vutils.save_image(inputImg_1.data,
            '%s/input_samples_iteration_%03d.png' % (opt.outf, iteration), normalize=True)



#-------------------train----------------------
for iteration in range(1,opt.niter+1):
    running_corrects = 0
    running_corrects_light = 0
    try:
        images_1,po_1,li_1,GT_1,by_image,same_po,same_li,same_id= data_train_1.next()
    except StopIteration:
        data_train_1 = iter(loader_train_1)
        images_1,po_1,li_1,GT_1,by_image,same_po,same_li,same_id = data_train_1.next()

    GT.data.resize_(GT_1.size()).copy_(GT_1)



    inputImg_1.data.resize_(images_1.size()).copy_(images_1)
    inputImg_2.data.resize_(by_image.size()).copy_(by_image)

    input_pose_1.data.resize_(po_1.size()).copy_(po_1)
    input_light_1.data.resize_(li_1.size()).copy_(li_1)

    same_pose.data.resize_(same_po.size()).copy_(same_po)
    same_light.data.resize_(same_li.size()).copy_(same_li)
    same_iden.data.resize_(same_id.size()).copy_(same_id)
    netS.zero_grad()

    output_pose_1, output_pose_2, output_light_1, output_light_2, out_f_1, out_f_2, out = netS(inputImg_1, inputImg_2)
    #-----------------mask test area-----------------------------
    # print out.data.type()
    # print GT.data.type()
    # print FinalMask.data.type() same
    # print FinalMask.data.size() 64x3x96x96
    # print FinalMask.data.size()
    # print out.data.size()

    Final_out = FinalMask * out

    Final_GT = FinalMask * GT


    #-----------------mask test area-----------------------------
    # f_1 & f_2 variable
    # same_iden variable
    err_recon = reconstructe_loss(Final_out, Final_GT)
    err_contraI = identity_contrastive_loss(out_f_1, out_f_2, same_iden)
    err_contraP = poss_contrastive_loss(output_pose_1, output_pose_2, same_pose)
    err_contraL = light_contrastive_loss(output_light_1,output_light_2, same_light)
    err_classL = light_class_loss(output_light_1, input_light_1)
    err_classP = pose_class_loss(output_pose_1, input_pose_1)
    # print(err_recon.data.size())
    # print(err_contraL.data.size())
    # print(err_classP.data.size())
    # modify the contrastive loss function to make contrastive loss be 1Lx1L 
    # contrastive loss and Softmax and Loss1 are all requires_grad
    # err_total = 1 * err_recon + 10 * err_contraP + 10 * err_contraI + 10 * err_classP + 20 * err_classL
    # err_total = err_recon + err_contraI + err_contraP + err_contraL + err_classL + err_classP
    err_total = w_r * err_recon
    err_total.backward()
    optimizerS.step()

    #----------------------Visualize-----------
    if(iteration % 200 == 0):

        pose_mtr.add(output_pose_1.data, input_pose_1.data)
        pose_trainacc = pose_mtr.value().diagonal().sum()*1.0/opt.batchSize
        pose_mtr.reset()

        light_mtr.add(output_light_1.data, input_light_1.data)
        light_trainacc = light_mtr.value().diagonal().sum()*1.0/opt.batchSize
        light_mtr.reset()
        #-----------------------------------------
        test(iteration, data_test, loader_test)
    # #pose prediction

    # preds_pose = torch.max(output_pose_1.data, 1)
    # running_corrects += torch.sum(preds == input_pose_1)
    # print('pose_accuracy: %.2f' 
    #         % (running_corrects * 1.0/images.size(0)))
    
    # #light prediction
    # preds_light = torch.max(output_light_1.data, 1)
    # running_corrects_light += torch.sum(preds_light == input_light_1)
    # print('light_accuracy: %.2f' 
    #         % (running_corrects_light * 1.0/images.size(0)))

    print('----------------------------------------')
    print('[%d/%d] Loss_S: %.4f ' %(iteration, opt.niter, err_total.data[0]))
    print('        Reco_S: %.4f ' %(err_recon.data[0]))
    print('        conL_S: %.4f ' %(err_contraL.data[0]))
    print('        conP_S: %.4f ' %(err_contraP.data[0]))
    print('        conI_S: %.4f ' %(err_contraI.data[0]))
    print('        Clas_P: %.4f ' %(err_classP.data[0]))
    print('        Clas_L: %.4f ' %(err_classL.data[0]))


    if(iteration % opt.save_step == 0):
        torch.save(netS.state_dict(), '%s/netS_%d.pth' % (opt.outf,iteration))


