<reponame>ZakZadeh/SelfGAN
"""
Saving SAM features
Developer: <NAME> 
"""

import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim as optim 
import torch.nn.functional as func
import model
import dataset
import utils

""" -------------------------------------------------------------------------"""
""" Parameters """
parser = argparse.ArgumentParser(description = 'Convolutional Neural Net')
parser.add_argument('--path', type=str, default='/data/zak/', help='dataset path')
parser.add_argument('--datasetName', type=str, default='kth', help='dataset name: mnist, f_mnist, cifar10, kth, ucf101, hmdb51, ballDrop3')
parser.add_argument('--startEpoch', type=int, default=100, help='number of epochs of pretrained model, default=0')
parser.add_argument('--useSpecNorm', type=bool, default=True, help='flag for using Spectral normalizatoin, default=False')
parser.add_argument('--useSelfAtten', type=bool, default=False, help='flag for using Self Attention Layer, default=False')
parser.add_argument('--nGPU', type=int, default=2, help='Num of GPUs available. Use 0 for CPU mode, default=0')
parser.add_argument('--nBatch', type=int, default=256, help='Batch size, default=16')
parser.add_argument('--nFeature', type=int, default=512, help='Feature size, default=16')
parser.add_argument('--imageSize', type=int, default=64, help='image dimension (weight & Height), default=64')
parser.add_argument('--nz', type=int, default=100, help='Size of z latent vector, default=100')
parser.add_argument('--ncf', type=int, default=64, help='Size of feature maps in critic, default=64')
parser.add_argument('--ngf', type=int, default=64, help='Size of feature maps in Generator, default=64')
parser.add_argument('--lrG', type=float, default=0.0001, help='optimizers Learning rate for Generator, default=0.0001')
parser.add_argument('--lrC', type=float, default=0.0004, help='optimizers Learning rate for critic, default=0.0004')
parser.add_argument('--beta1', type=float, default=0.5, help='beta1 hyperparameter for Adam optimizers, default=0.5')
parser.add_argument('--beta2', type=float, default=0.999, help='beta1 hyperparameter for Adam optimizers, default=0.999')
params = parser.parse_args()
params.nc = 3                 # Num of image channels
params.nClass = 10            # Num of classes
params.isInVideo = False      # Is Input video or still image
params.nFrames = 8            # num of video frames

""" -------------------------------------------------------------------------"""
""" initialization """

testLoader  = dataset.loadTest(params)
device = torch.device("cuda:0" if (torch.cuda.is_available() and params.nGPU > 0) else "cpu")

""" -------------------------------------------------------------------------"""
""" Loading Networks"""
netC  = model.Critic(params).to(device)
if (params.isInVideo):
    netG = model.Generator3(params).to(device)
    netC  = model.Critic3(params).to(device)
else:
    netG = model.Generator(params).to(device)
    netC  = model.Critic(params).to(device)
if (device.type == 'cuda') and (params.nGPU > 1):
    netG = nn.DataParallel(netG, list(range(params.nGPU)))
    netC = nn.DataParallel(netC, list(range(params.nGPU)))
    cudnn.benchmark = True

optimizerC = optim.Adam(netC.parameters(), lr = params.lrC, betas = (params.beta1, params.beta2))
optimizerG = optim.Adam(netG.parameters(), lr = params.lrG, betas = (params.beta1, params.beta2))

netG, netC, optimizerG, optimizerC = utils.loadCkpt(netG, netC, optimizerG, optimizerC, params)
net1 = nn.Sequential(*list(netC.module.children())[0:1])
net2 = nn.Sequential(*list(netC.module.children())[1:6])
maxpool = nn.MaxPool2d(4, stride=1)

# print(netC)
# print(net1)
# print(net2)

""" -------------------------------------------------------------------------"""
""" Getting Features """
print("Starting Getting Features...")
for fold in range(5):
    print("Fold " + str(fold))
    if params.isInVideo:
        pthMain = params.path + 'video/' + params.datasetName + '/features/'
    else:
        pthMain = params.path + 'image/' + params.datasetName + '/features/'

    utils.makeFolder(pthMain)
    pth = pthMain + str(fold)+'/'
    pthTrain = pth + 'train/'
    pthTest  = pth + 'test/'
    utils.makeFolder(pth)
    utils.makeFolder(pthTrain)
    utils.makeFolder(pthTest)

    labelTrain = torch.tensor(np.zeros((0)), dtype = torch.long, device = device)
    labelTest  = torch.tensor(np.zeros((0)), dtype = torch.long, device = device)

    for i, (data, label) in enumerate(testLoader, 0):
        if (i%5) != fold:
            inData = data.to(device)
            if (params.isInVideo):
                inData = dataset.sample(inData, params)
            batchSize = inData.size(0)
            if (batchSize != params.nBatch):
                continue
            label = label.to(device)
            
            if params.isInVideo:
                label = label - 1

            labelTrain = torch.cat((labelTrain, label), 0)
            output = net1(inData)
            output = torch.squeeze(output)
            output = net2(output)
            output = maxpool(output).squeeze()

            torch.save(output, pthTrain + 'train_{:04d}.pt'.format(i))
        else:
            if (params.isInVideo):
                inData = data.to(device)
                inData = dataset.sample(inData, params)
            else:
                inData  = data.to(device)

            batchSize = inData.size(0)
            if (batchSize != params.nBatch):
                continue
            label = label.to(device)
            if params.isInVideo:
                label = label - 1

            labelTest = torch.cat((labelTest, label), 0)

            output = net1(inData)
            output = torch.squeeze(output)
            output = net2(output)
            output = maxpool(output).squeeze()

            torch.save(output, pthTest + 'test_{:04d}.pt'.format(i))

    torch.save(labelTrain, pth + 'trainLabel.pt')
    torch.save(labelTest, pth + 'testLabel.pt')
