# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# + [markdown] id="view-in-github"
# <a href="https://colab.research.google.com/github/peterfazekas1999/-nofakes-project/blob/main/train_float_categorical_modified_11_01_2021.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + id="NTGJDkaxsrnw" colab={"base_uri": "https://localhost:8080/"} outputId="11a20a6e-0e12-4206-e99d-9ae35cb72029"
#code is used from these 3 repositories, have a look at them on GitHub or access the files from colab
# !git clone https://github.com/PeterWang512/FALdetector
# !git clone https://github.com/NVIDIA/flownet2-pytorch.git
# !git clone https://github.com/Kwanss/PCLNet

# + id="3W8h8ElvvMC7"
# -*- coding: utf-8 -*-
"""
Created on Wed Nov 25 10:32:09 2020

@author: <NAME>
"""
#import necessary modules and append paths 

import os
import csv
import pickle
import cv2
from PIL import Image
import matplotlib.pyplot as plt
import sys
sys.path.append("/content/FALdetector/networks/")
sys.path.append("/content/FALdetector/")
from drn_seg import DRNSeg
import torch
from torch.autograd import Variable
from torch.nn import Linear, ReLU, CrossEntropyLoss,BCELoss, Sequential, Conv2d, MaxPool2d, Module, Softmax, BatchNorm2d, Dropout
from torch.optim import Adam, SGD
sys.path.append("/content/flownet2-pytorch/")
import losses
from losses import MultiScale,EPE
sys.path.append("/content/PCLNet/Losses/")
sys.path.append("/content/PCLNet/models/")
from utils.tools import *
from utils.visualize import *
import pandas as pd
import numpy as np
import torch.nn as nn
# for reading and displaying images
from skimage.io import imread
# for creating validation set
from sklearn.model_selection import train_test_split
# for evaluating the model
from sklearn.metrics import accuracy_score
from tqdm import tqdm
from skimage.transform import rescale, resize, downscale_local_mean
from skimage import data, color
from torchsummary import summary

# + id="11kNl6zCuPPO" colab={"base_uri": "https://localhost:8080/"} outputId="87377767-612e-4fd6-b98a-5b509f6f38c2"
from google.colab import drive
drive.mount('/content/gdrive')
# !ln -s /content/gdrive/My\ Drive/ /mydrive
# !ls /mydrive
# !unzip /mydrive/nofakes/flow_pred_data/modified.zip -d modif
# !unzip /mydrive/nofakes/flow_pred_data/reference.zip -d ref
# !unzip /mydrive/nofakes/flow_pred_data/local_weight.zip

# + colab={"base_uri": "https://localhost:8080/"} id="RCc_Sot-srnw" outputId="44a23064-dfe2-4fd5-88dc-7bd52ca30de4"
pathM =r"/content/modif"
path =r"/content/ref"

#if you guys can come up with a more efficient way to do this go ahead.
height =400
width=400
#variable to control how many training examples to import
train_size =5
filenames = []

#Redo this part.... it doesnt import training images in order, I have printed out the filenames imported for convenience
def createTrain(path,n_images):
    arr=[]
    counter =0
    for r,d,f in os.walk(path):
        for _file in f:
          #make sure to change to correct image format .jpg or .png etc...
            if '.png' in _file:
                #print(_file)
                img = cv2.imread(r""+path+"/"+str(_file))
                filenames.append(path+"/"+str(_file))
                #keep image dimensions at 500 for now
                img = cv2.resize(img, (width,height), interpolation = cv2.INTER_AREA)
                #img = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
                arr.append(img)
                counter+=1
                #terminate once we imported the desired number of images
                if(counter==n_images):
                    return np.array(arr)
                    

X_ref=createTrain(path,train_size)
X_mod = createTrain(pathM,train_size)

print("the filenames are",filenames)
shape = X_ref.shape

# + colab={"base_uri": "https://localhost:8080/"} id="xn4RUUb4srnx" outputId="3a7bfb8d-aba9-4043-c78b-1404870c0945"
flow_arr =[]
#calcOptical gives an error if I input all training data at once (doesnt seem very efficient atm)
for i in range(shape[0]):
    flow = cv2.calcOpticalFlowFarneback(cv2.cvtColor(X_ref[i],cv2.COLOR_BGR2GRAY), cv2.cvtColor(X_mod[i],cv2.COLOR_BGR2GRAY),None, 0.5, 3, 4, 2, 3, 1.2, 0)
    flow_arr.append(flow)    
print(X_mod.shape)
flow_arr = np.array(flow_arr)
mag, ang = cv2.cartToPolar(flow[...,0], flow[...,1])
print(mag.shape)
path = "tester_before_train.png"
tester = save_heatmap_cv(X_mod[0], mag, path)
#print(flow)
print("flow arr has shape",flow_arr.shape)
#the code below is pretty similar to how the binary classifier is trained with a few changes 
#since the output shape is (2,height,width) for the vector field

# + [markdown] id="1WaL8ji9pVCE"
# **Discretize the flow fields**

# + colab={"base_uri": "https://localhost:8080/"} id="hKrXUh7npF0J" outputId="40bf8db8-d13c-4edd-b691-da9e5e386e92"
print("maximum flow value is: ",np.max(flow_arr))
print("minimum flow value is: ",np.min(flow_arr))
categorical_flow =[]
dic ={}
inv_dic={}
counter =0
#create a placeholder for the (u,v)->class pairs, hashmap is the easiest as its O(1) lookup time for the values
#Note when obtaining flow values from classes it takes O(n) time to look up each one, not very efficient so
#maybe re-do this part so that the hashmap key-value pairs are {class_pairs:(u,v)}

#max and minimum allowed flow values
max_f =5
min_f=-5
#fill up hashmap with values
for i in range(min_f,max_f+1):
  for j in range(min_f,max_f+1):
    dic[i,j]=counter
    inv_dic[counter]=[i,j]
    counter+=1
print(len(dic))
print("this is what the hashmap looks like",dic)
print("this is what the inverse looks like",inv_dic)
dimens= flow_arr.shape

#categorise the flow into distinct values
for flows in flow_arr:
  temp_flow =np.zeros((dimens[1],dimens[2]))
  print(temp_flow.shape)
  for i in range(dimens[1]):
    for j in range(dimens[2]):
      value = flows[i][j]
      #makes sure we are not going over the max and min values 
      if (value[1]>max_f or value[1]<min_f or value[0]>max_f or value[0]<min_f):
        continue
      temp_flow[i][j]=dic[int(value[0]),int(value[1])]
      
  categorical_flow.append(temp_flow)
categorical_flow = np.array(categorical_flow)
#print(categorical_flow.shape)
#print(categorical_flow[0][0])

# + id="1Vqlb8wk3mrb" colab={"base_uri": "https://localhost:8080/"} outputId="5db94ea7-0ae9-477d-ed70-1c69c8986d07"
#Sanity check to make sure shapes add up

print("ytrain has shape",categorical_flow.shape)
print("xtrain has shape",X_mod.shape)
#print(categorical_flow[0][100])

# + id="aphClU4NrJ3E"


# + [markdown] id="uM9U8BGOrJEv"
# **make sure flow is set up properly**

# + id="UKYz6MyxrHrx"


# + colab={"base_uri": "https://localhost:8080/"} id="z1WjTZ36yieZ" outputId="c3a97e6c-0a41-4571-f48e-9cd34d0f0be7"
print(torch.cuda.get_device_name(0))

# + id="1IGPecamsrnx" colab={"base_uri": "https://localhost:8080/"} outputId="f952f718-c440-41b8-c26f-afcf5bb2b372"
#check if GPU is available
cuda0 = torch.device('cuda:0')

if torch.cuda.is_available():
    device = 'cuda:{}'.format("cuda:0")
else:
    device = 'cpu'


model = DRNSeg(2)
device = torch.device(cuda0)
#state_dict = torch.load("local.pth", map_location=device)
#model.load_state_dict(state_dict['model'])
model.to(device)
#summary(model, (3, height, width))


  


# + id="tNP5FvjtW9A4" colab={"base_uri": "https://localhost:8080/"} outputId="8a250917-da8f-4e6a-f4e3-bf77f8c85c4c"
temp = torch.randn(1, 3, 200, 200)
temp = temp.cuda()
sizes =[2,4,16,32,64,121]
# Let's print it
shape=model(temp).shape
print(shape)

#leave it out for now as I am using the DRNSeg(121)
if False:
  perPix_model = nn.Sequential(model,nn.ReLU(),
                              nn.Conv2d(sizes[0],sizes[1],kernel_size=(1,1)).cuda(),nn.ReLU(),
                              nn.Conv2d(sizes[1],sizes[2],kernel_size=(1,1)).cuda(),nn.ReLU(),
                              nn.Conv2d(sizes[2],sizes[3],kernel_size=(1,1)).cuda(),nn.ReLU(),
                               nn.Conv2d(sizes[3],sizes[4],kernel_size=(1,1)).cuda(),nn.ReLU(),
                               nn.Conv2d(sizes[4],sizes[5],kernel_size=(1,1)).cuda(),
                              nn.Tanh()
                              )
#just add a tanh activation onto it
perPix_model = DRNSeg(len(dic)).to(device)
summary(perPix_model, (3, height, width))
perPix_model.to(device)

# + [markdown] id="jtVAZde-vW3w"
# **Training the Classifier**

# + colab={"base_uri": "https://localhost:8080/"} id="h1fXnALpsrnx" outputId="bc0f2495-467f-4f43-b4a9-1f4af706b103"
#changed crossentropy to MSE since cross entropy needs classes, we arent doing classification
#other losses might work much better 
from torch.utils.data import TensorDataset, DataLoader
optimizer = Adam(model.parameters(),lr=0.07) 
criterion =torch.nn.MSELoss()
criterion_classif = nn.CrossEntropyLoss()




#criterion = MultiScale(model)
if torch.cuda.is_available():
  model = model.cuda()
  perPix_model =perPix_model.cuda()
  criterion = criterion.cuda()
  criterion_classif=criterion_classif.cuda()
#convert to torch format
train_x = X_mod
train_x_orig=X_ref
train_x = train_x.reshape(train_size,3, height, width)
train_x_orig = train_x_orig.reshape(train_size,3, height, width)
print("train_x has shape",train_x.shape)
train_x  = torch.from_numpy(train_x)

#train_x_orig  = torch.from_numpy(train_x_orig).float()


# converting the target into torch format
train_y=flow_arr
cat_shape = categorical_flow.shape
categorical_label=torch.from_numpy(categorical_flow)
train_y = train_y.reshape(train_size,2,height,width)
print("train_y has shape: ",train_y.shape)
train_y = torch.from_numpy(train_y).float()
print("xtrain shape",train_x.shape)
print("y shape",categorical_label.shape)


#MAIN CHANGE, dataloader so that batches can be trained, use batch size of 1 and then multiple pictures can be trained
my_dataset = TensorDataset(train_x,categorical_label) # create your datset
b_size =1
my_dataloader = DataLoader(my_dataset,batch_size=b_size)



#this part of the code trains the classifier called perPix_model
def train_classif(epoch,b_size):
    correct =0
    #iterate over batches
    for idx,(x_t,y_t) in enumerate(my_dataloader):
      perPix_model.train()
      tr_loss = 0
      
      # getting the training set
      x_train, y_categorical = Variable(x_t), Variable(y_t)
      
      # converting the data into GPU format
      if torch.cuda.is_available():
          x_train = x_train.cuda()
          y_categorical = y_categorical.cuda()
      optimizer.zero_grad()
      
      # prediction for training and validation set
      output_train = perPix_model(x_train.float())
      
      
      #print("y has shape",y_categorical.shape)
      label_output = torch.argmax(output_train,dim=1)
      #print("label out shape",label_output.shape)
      #print("output shape",output_train.shape)
      Loss = criterion_classif(output_train, y_categorical.long())


      #print(label_output)
      #print("label output has shape",label_output.shape)
      correct = (label_output==y_categorical).sum().float()/(height*width*b_size)
      

      total = train_size
      
      #print(Loss)
      train_losses.append(Loss)
        
        # computing the updated weights of all the model parameters
      Loss.backward()
      optimizer.step()
      tr_loss = Loss.item()
      if epoch%10 == 0:
          # printing the validation loss
          print('Epoch : ',epoch+1, "for Batch:",idx, '\t', 'loss :',Loss.item(),"accuracy: ",correct.item())
      if epoch%1000 == 0:
        checkpoint_path = "gdrive/MyDrive/local_detector_model_{}.pt".format(int(epoch/1000))
        torch.save({
            'checkpoint_epoch': checkpoint_epoch+n_epochs_to_train,
            'model_state_dict': perPix_model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'train_losses': train_losses
            }, checkpoint_path)
      print(checkpoint_epoch+n_epochs_to_train)
    
        

# + colab={"base_uri": "https://localhost:8080/"} id="uD_FcDy6ueFx" outputId="ce34c7b1-ae23-4d3b-c622-dd8283562fb6"
loss = nn.CrossEntropyLoss()
input = torch.randn(3, 5, requires_grad=True)
target = torch.empty(3, dtype=torch.long).random_(5)
output = loss(input, target)
print(input.shape)
print(target.shape)

# + id="UUxCx2RD4s8n" colab={"base_uri": "https://localhost:8080/"} outputId="00f7a50a-7a9f-45ed-93b4-89f247af9aaf"
#Check for existing checkpoint in content/, train additional epochs if found
n_epochs_to_train = 4000
train_losses =[]
correct =0

#same parameters the FAL paper have used
optimizer = Adam(perPix_model.parameters(),lr=0.0001,betas =(0.9,0.999))

for i in reversed(range(100)):
  checkpoint_path = "gdrive/MyDrive/local_detector_model_{}.pt".format(i*1000)
  if os.path.exists(checkpoint_path):
      checkpoint = torch.load(checkpoint_path)
      perPix_model.load_state_dict(checkpoint['model_state_dict'])
      optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
      checkpoint_epoch = checkpoint['checkpoint_epoch']
      train_losses = checkpoint['train_losses']
      print("checkpoint loaded")
      print("checkpoint epoch:" + str(checkpoint_epoch))
      break
  if i==0:
    checkpoint_epoch=0
    print("No existing checkpoint found")


for i in range(n_epochs_to_train):
    train_classif(checkpoint_epoch+i,b_size)

print("In this run, another "+str(n_epochs_to_train)+" epochs were trained")
print("Total epochs model has been trained upon:" + str(checkpoint_epoch+n_epochs_to_train))
    

# + [markdown] id="moJiR05IJIv0"
# **plots the training loss**

# + id="fIXgHEy89vn4"
 print(len(train_losses))
iter = np.linspace(0,len(train_losses),len(train_losses))
plt.plot(iter,train_losses)
plt.show()


# + [markdown] id="7Ar-XxVCve6k"
# **Training the regression, skip these blocks for now**

# + id="d3oTou4fWELj"
MS_criterion= MultiScale(model)
EPE_l,MS_l,REC_l=1.5,15,1

#the function below is very messy atm, im just happy that it works, it needs to be cleaned up though
#it just converts and reshapes tensors to numpy arrays and vica versa during training for the recreational loss

#creates the recreation loss, unwarps image and compares with the original image
def rec_loss(output_train,x_train,x_train_orig):
    output_np = output_train.cpu()
    #print(output_np.shape)
    
    #get shapes of output and convert/reshape into np array
    with torch.no_grad():
      output_np = np.array(output_np)  
    shape = output_np.shape
    output_np =np.reshape(output_np,(shape[0],shape[2],shape[3],shape[1]))
    

    #get shapes of x_train and convert/reshape into np array
    x_train_np = x_train.detach().cpu().numpy()
    shape2 =x_train_np.shape
    x_train_np =np.reshape(x_train_np,(shape2[0],shape2[2],shape2[3],shape2[1]))


    #print(x_train_np.shape,type(x_train_np))
    x_warped=[]
    #undo the warps using prediction flow field and manipulated image
    for i in range(train_size):
      x_warped.append(warp(x_train_np[i],output_np[i]))
    #convert from torch to numpy and reshape
    x_warped = np.array(x_warped)
    shape_warp = x_warped.shape

    #get unwarped image and turn it into torch format
    x_warped = np.reshape(x_warped,(shape_warp[0],shape_warp[3],shape_warp[1],shape_warp[2]))
    x_warp_tens = torch.from_numpy(x_warped).float()
    #turn it into cuda tensor
    if torch.cuda.is_available():
      x_warp_tens = x_warp_tens.cuda()
    #return the norm of the warped image-original image as the loss
    return torch.norm(x_warp_tens-x_train_orig,p=2,dim=1).mean()



def train_regress(epoch,EPE_l,MS_l,REC_l):
    model.train()
    tr_loss = 0
    # getting the training set
    x_train, y_train = Variable(train_x), Variable(train_y)
    x_train_orig=Variable(train_x_orig)
    # converting the data into GPU format
    if torch.cuda.is_available():
        x_train = x_train.cuda()
        y_train = y_train.cuda()
        x_train_orig=x_train_orig.cuda()  
    # clearing the Gradients of the model parameters
    optimizer.zero_grad()
    # prediction for training and validation set
    output_train = model(x_train)
    print("output train shape",output_train.shape)
    print(output_train)
    total = train_size
    # computing the training and validation loss
    #print(x_train.shape)
    #print(output_train.shape)
    print("y train shape",y_train.shape)
    print(y_train)
    REC_loss = rec_loss(output_train,x_train,x_train_orig)
    EPE_loss_train = EPE(output_train,y_train.float())
    MS_loss_train,_ = MS_criterion(output_train, y_train.float())
    
    loss_train=EPE_l*EPE_loss_train+MS_l*MS_loss_train +REC_l*REC_loss 
    train_losses.append(loss_train)
    
    # computing the updated weights of all the model parameters
    loss_train.backward()
    optimizer.step()
    tr_loss = loss_train.item()

    if epoch%10 == 0:
        # printing the validation loss
        print('Epoch : ',epoch+1, '\t', 'loss :',loss_train.item(),"accuracy: ")
        print("epe: ",EPE_loss_train.item()," MS: ",MS_loss_train.item()," reconstruct: ",REC_loss.item())
        print("\n")


# + id="iy-SmfT1srny"
n_epochs = 500
train_losses =[]
correct =0
#same parameters the FAL paper have used
optimizer = Adam(model.parameters(),lr=0.0001,betas =(0.9,0.999)) 
for epoch in range(n_epochs): 
    train_regress(epoch,EPE_l,MS_l,REC_l)


# + id="Bn7gOX8Usrny"
#function to
def get_heatmap_cv(img, magn, max_flow_mag):
    min_flow_mag = .5
    cv_magn = np.clip(
        255 * (magn - min_flow_mag) / (max_flow_mag - min_flow_mag),
        a_min=0,
        a_max=255).astype(np.uint8)
    if img.dtype != np.uint8:
        img = (255 * img).astype(np.uint8)
        
    heatmap_img = cv2.applyColorMap(cv_magn, cv2.COLORMAP_MAGMA)
    heatmap_img = heatmap_img[..., ::-1]
    h, w = magn.shape
    img_alpha = np.ones((h, w), dtype=np.double)[:, :, None]
    heatmap_alpha = np.clip(
        magn / max_flow_mag, a_min=0, a_max=1)[:, :, None]**.7
    heatmap_alpha[heatmap_alpha < .2]**.5
    pm_hm = heatmap_img * heatmap_alpha
    pm_img = img * img_alpha
    cv_out = pm_hm + pm_img * (1 - heatmap_alpha)
    cv_out = np.clip(cv_out, a_min=0, a_max=255).astype(np.uint8)
    return cv_out

def save_heatmap_cv(img, magn, path, max_flow_mag=2):
    cv_out = get_heatmap_cv(img, magn, max_flow_mag)
    out = Image.fromarray(cv_out)
    plt.imshow(out)
    plt.show()
    out.save(path, quality=95)
    return out


# + id="zQOqmEbmsrny"
from utils.tools import *
from utils.visualize import *


import torchvision.transforms as transforms
tf = transforms.Compose([transforms.ToTensor(),
                         transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                              std=[0.229, 0.224, 0.225])])

# + id="eAdi8lECsrny"
#quick sanity check again

#print(flow.shape)
#print(pic.size)

# + [markdown] id="DM2lLcMnIcha"
# **uses training image to predict the flow field**

# + id="1wg4-nudsrny"
test_ref = X_ref[2]
test_mod = X_mod[2]
ground_flow = cv2.calcOpticalFlowFarneback(cv2.cvtColor(test_ref,cv2.COLOR_BGR2GRAY), cv2.cvtColor(test_mod,cv2.COLOR_BGR2GRAY),None, 0.5, 3, 4, 2, 3, 1.2, 0)
image = np.array([test_mod])

shapes = image.shape

image = np.reshape(image,(shapes[0],shapes[3],shapes[1],shapes[2]))
image_tensor = torch.from_numpy(image)
image_tensor= Variable(image_tensor)
# converting the data into GPU format
if torch.cuda.is_available():
    image_tensor = image_tensor.cuda()
with torch.no_grad():
        flow = perPix_model(image_tensor.float()).float().cpu()
        print(flow.shape)
        flow = torch.argmax(flow,dim =1).numpy()
        flow = np.reshape(flow,(shapes[1],shapes[2]))
        
        print(flow)
        real_flow = np.zeros((2,400,400))
        temp =0
        for i in range(shapes[1]):
          for j in range(shapes[2]):
            real_flow[:,i,j]=list(dic.keys())[list(dic.values()).index(flow[i][j])]
        print(real_flow.shape)
        np.set_printoptions(edgeitems=3)
        #print(real_flow)
        real_flow = np.transpose(real_flow, (1, 2, 0))
        print(real_flow.shape)
        #print(real_flow)
        w, h, _ = real_flow.shape

# + id="NH_Cwm4Psrny"
image = np.reshape(image,(h,w,3))
image = Image.fromarray(image,"RGB")
modified = image.resize((h, w),Image.BICUBIC)
modified_np = np.asarray(modified)
flow = flow_resize(real_flow, modified.size)
#print(flow.shape)
mag, ang = cv2.cartToPolar(flow[...,1], flow[...,0])
#print(mag.shape)
#print(modified_np.shape)
#save your image to the folder
path = "tester.png"
tester =get_heatmap_cv(modified_np,mag,max_flow_mag=10)
plt.title("Predicted modification")
plt.imshow(tester)

# + [markdown] id="RI3Yagn9IznO"
# **plots the ground truth heatmap on the image**

# + id="0AbiG2ez9tzm"

mag, ang = cv2.cartToPolar(ground_flow[...,0], ground_flow[...,1])
#print(mag.shape)
real = get_heatmap_cv(modified_np,mag,max_flow_mag=12)
plt.title("Ground truth modification")
plt.imshow(real)

# + id="26P3Me_CGm9H"
~

StarcoderdataJupyterScriptsDedupFiltered