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

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/antoniomuso/speech2face/blob/master/face_processing.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + id="8BQpw02EO21K" colab_type="code" colab={}
from google.colab import drive
drive.mount('/content/drive')

# + id="omXh_0lnavT6" colab_type="code" colab={}
# ! cp "/content/drive/My Drive/Speech2Face/zippedFaces.tar.gz" /content
# ! tar zxf zippedFaces.tar.gz
# ! rm zippedFaces.tar.gz

# + id="w59EDQVO8r6v" colab_type="code" colab={}
# #! pip install openface
# ! pip3 install face_recognition
from face_recognition import face_landmarks, face_locations, face_encodings

# + id="s9uk5jzR8x7j" colab_type="code" colab={}
# #! pip install deepface
#from deepface import DeepFace

# + id="76KnWmm09EJ0" colab_type="code" colab={}
# Workspace cleanup

# ! rm -r Face_Feature/
# ! rm -r Faces_rearranged/
# ! rm face_features*.zip 

# + id="AOEtoMbdNYWA" colab_type="code" colab={}
import glob

import numpy as np
import sys

from PIL import Image
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
import  torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from random import randint

import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import optim

import os
from os.path import join
from os import listdir

from tqdm import tqdm
from google.colab.patches import cv2_imshow
import cv2
import sys

import random

# + id="TmqI6zWd9dLR" colab_type="code" colab={}
import os, sys
from IPython.utils import io

class HiddenPrints:
    def __enter__(self):
        self._original_stdout = sys.stdout
        sys.stdout = open(os.devnull, 'w')

    def __exit__(self, exc_type, exc_val, exc_tb):
        sys.stdout.close()
        sys.stdout = self._original_stdout


# + id="hxf2WEJFNpKF" colab_type="code" colab={}
class Vgg_face_dag(nn.Module):

    def __init__(self):
        super(Vgg_face_dag, self).__init__()
        self.meta = {'mean': [129.186279296875, 104.76238250732422, 93.59396362304688],
                     'std': [1, 1, 1],
                     'imageSize': [224, 224, 3]}
        self.conv1_1 = nn.Conv2d(3, 64, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu1_1 = nn.ReLU(inplace=True)
        self.conv1_2 = nn.Conv2d(64, 64, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu1_2 = nn.ReLU(inplace=True)
        self.pool1 = nn.MaxPool2d(kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False)
        self.conv2_1 = nn.Conv2d(64, 128, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu2_1 = nn.ReLU(inplace=True)
        self.conv2_2 = nn.Conv2d(128, 128, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu2_2 = nn.ReLU(inplace=True)
        self.pool2 = nn.MaxPool2d(kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False)
        self.conv3_1 = nn.Conv2d(128, 256, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu3_1 = nn.ReLU(inplace=True)
        self.conv3_2 = nn.Conv2d(256, 256, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu3_2 = nn.ReLU(inplace=True)
        self.conv3_3 = nn.Conv2d(256, 256, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu3_3 = nn.ReLU(inplace=True)
        self.pool3 = nn.MaxPool2d(kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False)
        self.conv4_1 = nn.Conv2d(256, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu4_1 = nn.ReLU(inplace=True)
        self.conv4_2 = nn.Conv2d(512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu4_2 = nn.ReLU(inplace=True)
        self.conv4_3 = nn.Conv2d(512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu4_3 = nn.ReLU(inplace=True)
        self.pool4 = nn.MaxPool2d(kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False)
        self.conv5_1 = nn.Conv2d(512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu5_1 = nn.ReLU(inplace=True)
        self.conv5_2 = nn.Conv2d(512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu5_2 = nn.ReLU(inplace=True)
        self.conv5_3 = nn.Conv2d(512, 512, kernel_size=[3, 3], stride=(1, 1), padding=(1, 1))
        self.relu5_3 = nn.ReLU(inplace=True)
        self.pool5 = nn.MaxPool2d(kernel_size=[2, 2], stride=[2, 2], padding=0, dilation=1, ceil_mode=False)
        self.fc6 = nn.Linear(in_features=25088, out_features=4096, bias=True)
        self.relu6 = nn.ReLU(inplace=True)
        self.dropout6 = nn.Dropout(p=0.5)
        self.fc7 = nn.Linear(in_features=4096, out_features=4096, bias=True)
        self.relu7 = nn.ReLU(inplace=True)
        self.dropout7 = nn.Dropout(p=0.5)
        self.fc8 = nn.Linear(in_features=4096, out_features=2622, bias=True)

    def forward(self, x0, is_fc8=False):
        x1 = self.conv1_1(x0)
        x2 = self.relu1_1(x1)
        x3 = self.conv1_2(x2)
        x4 = self.relu1_2(x3)
        x5 = self.pool1(x4)
        x6 = self.conv2_1(x5)
        x7 = self.relu2_1(x6)
        x8 = self.conv2_2(x7)
        x9 = self.relu2_2(x8)
        x10 = self.pool2(x9)
        x11 = self.conv3_1(x10)
        x12 = self.relu3_1(x11)
        x13 = self.conv3_2(x12)
        x14 = self.relu3_2(x13)
        x15 = self.conv3_3(x14)
        x16 = self.relu3_3(x15)
        x17 = self.pool3(x16)
        x18 = self.conv4_1(x17)
        x19 = self.relu4_1(x18)
        x20 = self.conv4_2(x19)
        x21 = self.relu4_2(x20)
        x22 = self.conv4_3(x21)
        x23 = self.relu4_3(x22)
        x24 = self.pool4(x23)
        x25 = self.conv5_1(x24)
        x26 = self.relu5_1(x25)
        x27 = self.conv5_2(x26)
        x28 = self.relu5_2(x27)
        x29 = self.conv5_3(x28)
        x30 = self.relu5_3(x29)
        x31_preflatten = self.pool5(x30)
        x31 = x31_preflatten.view(x31_preflatten.size(0), -1)
        x32 = self.fc6(x31)
        x33 = self.relu6(x32)
        x34 = self.dropout6(x33)
        x35 = self.fc7(x34)
        x36 = self.relu7(x35)
        x37 = self.dropout7(x36)
        if is_fc8:
            x38 = self.fc8(x37)
        else:
            x38 = x37
        return x38


def vgg_face_dag(weights_path=None, **kwargs):
    """
    load imported model instance
    Args:
        weights_path (str): If set, loads model weights from the given path
    """
    model = Vgg_face_dag()
    if weights_path:
        state_dict = torch.load(weights_path)
        model.load_state_dict(state_dict)
    return model


# + id="HIkwNAYcPQri" colab_type="code" colab={}
class _Normalize_Tensor(object):
    def __init__(self, color_mean, color_std):
        self.color_mean = color_mean
        self.color_std = color_std

    def __call__(self, img):
        # Convert image to Tensor
        img = transforms.functional.to_tensor(img)

        # Normalize image by the parameter of pre-trained face-encoder
        img = transforms.functional.normalize(
            img, self.color_mean, self.color_std)

        return img


# + id="MMEt6gSvOvSr" colab_type="code" colab={}
datasetPath = 'unzippedFaces'
base_dir='/content/drive/My Drive/Speech2Face/'
data_points = os.listdir(datasetPath) ## GETTING DATA POINTS THAT IS NAME OF CELEBRTIES.

# + id="RPcvoH0zN7mH" colab_type="code" colab={}
device = 'cuda'
face_encoder_path = join(base_dir, "Pretrained", "vgg_face_dag.pth")
face_encoder_model = vgg_face_dag(face_encoder_path)
face_encoder_model.eval()
face_encoder_model.to(device)

color_mean, color_std = face_encoder_model.meta["mean"], face_encoder_model.meta["std"]
color_mean = [tmp / 255.0 for tmp in color_mean]
color_std = [tmp / 255.0 for tmp in color_std]

transform_fe = transforms.Compose([
            transforms.Resize((224, 224)),
            _Normalize_Tensor(color_mean, color_std)
])



# + id="rH6sRJBg9vvp" colab_type="code" colab={}
# Deep-head-pose-LITE cloning
# ! git clone https://github.com/OverEuro/deep-head-pose-lite.git
# ! mv deep-head-pose-lite dhpl

# + id="lpMS9Xcps3GQ" colab_type="code" colab={}
from dhpl import stable_hopenetlite
pos_net = stable_hopenetlite.shufflenet_v2_x1_0()
saved_state_dict = torch.load('dhpl/model/shuff_epoch_120.pkl', map_location="cpu")
pos_net.load_state_dict(saved_state_dict, strict=False)
pos_net.eval()

# + id="D1Qh9YEyExhC" colab_type="code" colab={}
# Deep-head-pose cloning
# ! git clone https://github.com/natanielruiz/deep-head-pose.git
# ! mv deep-head-pose/ dhp
# ! cp /content/drive/My\ Drive/Speech2Face/hopenet_robust_alpha1.pkl .

# + id="izNfNfCAF0fm" colab_type="code" colab={}
from dhp.code import hopenet
import torchvision

saved_state_dict = torch.load('/content/hopenet_robust_alpha1.pkl', map_location="cuda")
#print(saved_state_dict.keys())
hopenet_model = hopenet.Hopenet(torchvision.models.resnet.Bottleneck, [3, 4, 6, 3], 66)
hopenet_model.load_state_dict(saved_state_dict, strict=False)
hopenet_model.to(device)
hopenet_model.eval()

pos_net = hopenet_model


# + id="qYpu79SEvl70" colab_type="code" colab={}
def softmax_temperature(tensor, temperature):
    result = torch.exp(tensor / temperature)
    result = torch.div(result, torch.sum(result, 1).unsqueeze(1).expand_as(result))
    return result


transformations = transforms.Compose([transforms.Scale(224),
    transforms.CenterCrop(224), transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])

def y_p_r(img):
  input = transformations(img).to(device)
  yaw, pitch, roll = pos_net(input.unsqueeze(0))
  yaw, pitch, roll = (yaw.cpu().detach(), pitch.cpu().detach(), roll.cpu().detach())
  idx_tensor = torch.arange(66,dtype=torch.float)
  # _, yaw_bpred = torch.max(yaw.data, 1)
  # _, pitch_bpred = torch.max(pitch.data, 1)
  # _, roll_bpred = torch.max(roll.data, 1)
  # Continuous predictions
  yaw_predicted = softmax_temperature(yaw.data, 1)
  pitch_predicted = softmax_temperature(pitch.data, 1)
  roll_predicted = softmax_temperature(roll.data, 1)
  yaw_predicted = torch.sum(yaw_predicted * idx_tensor, 1).cpu() * 3 - 99
  pitch_predicted = torch.sum(pitch_predicted * idx_tensor, 1).cpu() * 3 - 99
  roll_predicted = torch.sum(roll_predicted * idx_tensor, 1).cpu() * 3 - 99

  return yaw_predicted, pitch_predicted, roll_predicted


#print(yaw_predicted, pitch_predicted, roll_predicted)

# + id="JIytMAqvPKJT" colab_type="code" colab={}
no_of_record_per_image = 40 
face_feature = np.zeros((len(data_points)*no_of_record_per_image,4096)) ### ARRAY THAT WILL STORE THE FACE FEATURE
face_feature_fc8 = np.zeros((len(data_points)*no_of_record_per_image,2622)) ### ARRAY THAT WILL STORE THE FACE FEATURE (fc8 layer)

if not("Face_Feature" in os.listdir()): 
    os.mkdir("Face_Feature")
    os.mkdir("Face_Feature/Faces")
index = 0
### EXTRACTING FACE FEATURE.
for i in tqdm(range(len(data_points))):
    for j in range(no_of_record_per_image):
      try: 
        path_img = join(datasetPath,data_points[i],"1.6",random.choice(listdir(join(datasetPath,data_points[i],"1.6"))))
        image_path = join(path_img ,  random.choice(listdir(path_img)))
        img = Image.open(image_path)
      except:
        continue

      pos = face_locations(np.array(img), model='cnn')

      if len(pos) == 0:
       continue   

      y, p, r = y_p_r(img)

      #display(img)
      #print(y, p, r)

      if not (-10 < y < 10) or not (-15 <= p <= 15):
        continue

      encoding = face_encoder_model(transform_fe(img).unsqueeze_(0).to(device))
      encoding_fc8 = face_encoder_model(transform_fe(img).unsqueeze_(0).to(device), is_fc8=True)
      # face_feature[index] = encoding.cpu().detach().numpy()
      face_feature[index,:] = encoding.cpu().detach().numpy()
      face_feature_fc8[index,:] = encoding_fc8.cpu().detach().numpy()

      # print(face_feature[index].sum())
      filename = data_points[i] + "_"  + str(index) + ".jpg"
      path_save = join("Face_Feature/Faces", filename)

      img.save(path_save)
    
      index += 1
      
    
np.save("Face_Feature/facefeature.npy",face_feature[:index])
np.save("Face_Feature/facefeature_fc8.npy",face_feature_fc8[:index])
print(index)

# + id="9LoAxZjVRZyv" colab_type="code" colab={}
# #!rm -r filtered_unzippedFaces/

PICS_PER_ACTOR = 40
face_feature = np.zeros((len(data_points)*PICS_PER_ACTOR,4096)) ### ARRAY THAT WILL STORE THE FACE FEATURE
face_feature_fc8 = np.zeros((len(data_points)*PICS_PER_ACTOR,2622)) ### ARRAY THAT WILL STORE THE FACE FEATURE (fc8 layer)

index = 0

face_dir = "Face_Feature"
faces_dir = "Face_Feature/Faces"
if not os.path.exists(face_dir):
  os.makedirs(faces_dir)


for i in tqdm(range(len(data_points))):
  dir = data_points[i]

  s = join(datasetPath, dir+"/1.6/")

  videos = listdir(s)

  #pics_per_video = PICS_PER_ACTOR / len(videos)

  taken_per_actor = 0

  for subdir in videos:
    t = join(s, subdir)

    #taken_per_video = 0
    if taken_per_actor >= PICS_PER_ACTOR:
      break

    for image in listdir(t):
      if taken_per_actor >= PICS_PER_ACTOR: #or taken_per_video >= pics_per_video:
        break

      with torch.no_grad():

        image_path = join(t, image)
        img = Image.open(image_path)
  
        pos = face_locations(np.array(img), model='cnn')
  
        if len(pos) == 0:
         continue   
  
        y, p, r = y_p_r(img)
  
        if not (-10 < y < 10) or not (-15 <= p <= 15):
          continue
  
        taken_per_actor += 1
        #taken_per_video += 1
  
        encoding = face_encoder_model(transform_fe(img).unsqueeze_(0).to(device))
        encoding_fc8 = face_encoder_model(transform_fe(img).unsqueeze_(0).to(device), is_fc8=True)
        # face_feature[index] = encoding.cpu().detach().numpy()
        face_feature[index,:] = encoding.cpu().detach().numpy()
        face_feature_fc8[index,:] = encoding_fc8.cpu().detach().numpy()
  
        # print(face_feature[index].sum())
        filename = dir + "_"  + str(index) + ".jpg"
  
        out_file = join(faces_dir, filename)
  
        img.save(out_file)
        index += 1  
      
np.save("Face_Feature/facefeature.npy",face_feature[:index])
np.save("Face_Feature/facefeature_fc8.npy",face_feature_fc8[:index])

# + id="W3dqs5hZUH31" colab_type="code" colab={}
 print(face_feature[1].sum())
 print(face_feature.shape)

# + id="-JT3mR-raCFz" colab_type="code" colab={}
# ! zip -r face_features_40_per_actor.zip filtered_unzippedFaces/ 
# ! cp face_features_40_per_actor.zip /content/drive/My\ Drive/Speech2Face/

# + id="9QjFbpSKQWQR" colab_type="code" colab={}
import shutil

images_path = '/content/Face_Feature/Faces'
out_path = '/content/Faces_rearranged'

def rearrange_pictures(): # Rearrange pictures following the dataset's structure
  for image in os.listdir(images_path):
    full_path = os.path.join(images_path, image)
    actor_name = image[:image.rindex('_')]
  
    actor_folder = os.path.join(out_path, actor_name)
    #print(actor_folder)
    
    if not os.path.exists(actor_folder):
      os.makedirs(actor_folder)
    shutil.move(full_path, actor_folder)

rearrange_pictures()

StarcoderdataJupyterScriptsDedupFiltered