<jupyter_start><jupyter_text>set up<jupyter_code>from google.colab import drive
drive.mount('/content/drive')
!pip install facenet_pytorch
!pip install lpips
!pip install pytorch-msssim
!pip install wandb
import os
!git clone https://github.com/danielroich/Face-Identity-Disentanglement-via-StyleGan2.git 'project'
CODE_DIR = 'project'
os.chdir(f'./{CODE_DIR}')
BASE_PATH = '/content/drive/MyDrive/CNN-project-weights/'
MOBILE_FACE_NET_WEIGHTS_PATH = BASE_PATH + 'CNN-project-weights/mobilefacenet_model_best.pth.tar'
GENERATOR_WEIGHTS_PATH = BASE_PATH + 'CNN-project-weights/550000.pt'
E_ID_LOSS_PATH = BASE_PATH + 'CNN-project-weights/model_ir_se50.pth'
IMAGE_DATA_DIR = BASE_PATH + 'Dataset/small_image/'
W_DATA_DIR = BASE_PATH + 'Dataset/small_w/'
MODELS_DIR = BASE_PATH + 'Models/'
from Configs import Global_Config
from Configs.training_config import config, GENERATOR_IMAGE_SIZE
from Training.trainer import Trainer
from torch.utils.data import DataLoader, random_split
from Models.Encoders.Landmark_Encoder import Landmark_Encoder
from Models.Encoders.ID_Encoder import ID_Encoder
from Models.Encoders.Inception import Inception
from Models.Discrimanator import Discriminator
from Models.LatentMapper import LatentMapper
from Models.StyleGan2.model import Generator
from Utils.data_utils import get_w_image, Image_W_Dataset, cycle_images_to_create_diff_order
import time
import torch
import torch.utils.data
from tqdm import tqdm
from Losses import id_loss<jupyter_output><empty_output><jupyter_text>network<jupyter_code>id_encoder = id_loss.IDLoss(E_ID_LOSS_PATH)
attr_encoder = Inception()
discriminator = Discriminator()
mlp = LatentMapper()
landmark_encoder = Landmark_Encoder.Encoder_Landmarks(MOBILE_FACE_NET_WEIGHTS_PATH)
generator = Generator(GENERATOR_IMAGE_SIZE, 512, 8)
state_dict = torch.load(GENERATOR_WEIGHTS_PATH)
generator.load_state_dict(state_dict['g_ema'], strict=False)
# load our checkpoints

attr_encoder = torch.load(MODELS_DIR + 'attr_encoder_HPMNYYOTTIXN_1616315604.9508004_0.pt')
mlp = torch.load(MODELS_DIR + 'maper_HPMNYYOTTIXN_1616315604.542634_0.pt')
id_encoder = id_encoder.to(Global_Config.device)
attr_encoder = attr_encoder.to(Global_Config.device)
discriminator = discriminator.to(Global_Config.device)
mlp = mlp.to(Global_Config.device)
generator = generator.to(Global_Config.device)
landmark_encoder = landmark_encoder.to(Global_Config.device)
id_encoder = id_encoder.eval()
attr_encoder = attr_encoder.eval()
discriminator = discriminator.eval()
generator = generator.eval()
mlp = mlp.eval()
landmark_encoder = landmark_encoder.eval()<jupyter_output><empty_output><jupyter_text>get images<jupyter_code>w_image_dataset = Image_W_Dataset(W_DATA_DIR, IMAGE_DATA_DIR)
data_loader = DataLoader(dataset=w_image_dataset, batch_size=config['batchSize'], shuffle=False)
data = next(iter(data_loader))
ws, images = data
import matplotlib.pyplot as plt
plot_images = images.cpu().detach().numpy().transpose(0, 2, 3, 1)
# choose idx 
id_idx = 5
attr_idx = (config['batchSize']-1) if id_idx==0 else (id_idx - 1)
# identity

plt.axis('off')
plt.imshow(plot_images[id_idx])
plt.show()
# attribute

plt.axis('off')
plt.imshow(plot_images[attr_idx])
plt.show()<jupyter_output><empty_output><jupyter_text>have fun!<jupyter_code>test_ws = ws.to(Global_Config.device)
test_id_images = images.to(Global_Config.device)
test_attr_images_cycled = cycle_images_to_create_diff_order(test_id_images)
def get_concat_vec(id_images, attr_images, id_encoder, attr_encoder):
    with torch.no_grad():
        id_vec = torch.squeeze(id_encoder.extract_feats((id_images * 2) - 1))
        attr_vec = torch.squeeze(attr_encoder(attr_images))
        test_vec = torch.cat((id_vec, attr_vec), dim=1)
        return test_vec
concat_vec_cycled = get_concat_vec(test_id_images, test_attr_images_cycled, id_encoder, attr_encoder)
with torch.no_grad():
  mapped_concat_vec_cycled = mlp(concat_vec_cycled)
with torch.no_grad():
  cycled_generated_image = get_w_image(mapped_concat_vec_cycled[id_idx], generator)
# Result!

plt.axis('off')
plt.imshow(cycled_generated_image)
plt.show()<jupyter_output>Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).
{'fraction_non_alphanumeric': 0.10256410256410256, 'fraction_numerical': 0.01566951566951567, 'mean_word_length': 7.650924024640657, 'pattern_counts': {'":': 0, '<': 16, '<?xml version=': 0, '>': 16, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '607391', 'license': 'MIT', 'n_tokens_mistral': 1580, 'n_tokens_neox': 1513, 'n_words': 260, 'path': 'Notebooks/Inference.ipynb', 'repo_name': 'YuGong123/ID-disentanglement-Pytorch', 'size': '4212'}