<jupyter_start><jupyter_text>Latent-Based Regression through GAN SemanticsThis notebook contains an example playground for the paper: LARGE: Latent-Based Regression through GAN SemanticsWe currently support both sorting and regression models. Expand the relevant section for more info.<jupyter_code>#@title Setup required packages, models. This may take a few minutes.

import os

data_dir        = os.path.join("/content", "latent_regression")
celeba_dir      = os.path.join(data_dir, "celeba_images")
latents_dir     = os.path.join(data_dir, "latents")
cat_dir         = os.path.join(data_dir, "afhq_cat_images")
cat_latents_dir = os.path.join(data_dir, "cat_latents")

# install requirements
!pip install torch==1.7.1+cu110 torchvision==0.8.2+cu110 torchaudio==0.7.2 -f https://download.pytorch.org/whl/torch_stable.html 

!git clone https://github.com/omertov/encoder4editing.git
!wget https://github.com/ninja-build/ninja/releases/download/v1.8.2/ninja-linux.zip
!sudo unzip ninja-linux.zip -d /usr/local/bin/
!sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force

!pip install ftfy regex tqdm 
!pip install git+https://github.com/openai/CLIP.git

!git clone https://github.com/YotamNitzan/LARGE.git

!gdown --id 1WqhoitsRZ2nfUslIMf5Rg_WWOp2LW2Hu

!unzip -q latent_regression.zip
!unzip -q latent_regression/celeba_images.zip -d $celeba_dir
!unzip -q latent_regression/latents.zip -d $latents_dir
!unzip -q latent_regression/afhq_cat_images -d $cat_dir
!unzip -q latent_regression/afhq_cat_latents -d $cat_latents_dir

%load_ext autoreload
%autoreload 2

#@title General imports and variables
from PIL import Image

import sys

import numpy as np

set_feature_directions = {"celebA": "precomputed_feature_dirs.npy",
                          "afhq_cat": "precomputed_cat_feature_dirs.npy"}

set_num_layers = {"celebA": 18,
                  "afhq_cat": 16}<jupyter_output><empty_output><jupyter_text>Sorting Sorting in broken into two parts:1.   Choose a semantic direction. Currently supported options are:    *   Choose a direction from the paper    *   Create your own using a StyleCLIP-based approach    *   Upload a boundary file2.   Run the sorting script.  1) Choose a latent direction:  Precomputed paper directions:<jupyter_code>#@title Precomputed boundaries:

boundary = 'celebA_hair_color' #@param['celebA_hair_length', 'celebA_hair_color', 'celebA_makeup', 'celebA_expression', 'cat_age', 'cat_yaw', 'cat_pitch']

dataset = boundary.split("_")[0]
if dataset == "cat":
    dataset = "afhq_cat"

clip_boundary_dir = os.path.join(data_dir, "clip_boundaries")
sefa_boundary_dir = os.path.join(data_dir, "sefa_boundaries")

boundary_type = {"celebA": "s", "afhq_cat": "w"}[dataset]
latent_src = "e4e"
latent_type = {"celebA": "s_latents", "afhq_cat": "latents"}[dataset]
boundary_ext = {"celebA": ".pickle", "afhq_cat": ".npy"}[dataset]

boundary_path = {"celebA_hair_length": os.path.join(clip_boundary_dir, "short_to_long_S_10"),
                 "celebA_hair_color":  os.path.join(clip_boundary_dir, "black_to_blonde_S"),
                 "celebA_makeup":      os.path.join(clip_boundary_dir, "makeup_S"),
                 "celebA_expression":  os.path.join(clip_boundary_dir, "sad_to_smile_S"),
                 "cat_age":            os.path.join(sefa_boundary_dir, "cat_age"),
                 "cat_yaw":            os.path.join(sefa_boundary_dir, "cat_yaw"),
                 "cat_pitch":          os.path.join(sefa_boundary_dir, "cat_pitch")}[boundary]

boundary_path = os.path.join(boundary_path, f'boundary{boundary_ext}')<jupyter_output><empty_output><jupyter_text>Create your own direction with StyleCLIP:*   Enter source and target texts which describe the sorting direction.*   Choose a cutoff value for considered directions. A higher value gives a more disentangled direction, at the risk of removing too much relevant information.<jupyter_code>#@title Generate a StyleCLIP boundary:

dataset = 'afhq_cat' #@param['celebA', 'afhq_cat']

source_text = "white fur" #@param {type: "string"}
target_text = "black fur" #@param {type: "string"}

cutoff_percentile = 80 #@param {type: "integer"}

boundary_type = "w+"
latent_src = "psp"
latent_type = "latents"
boundary_ext = ".npy"
clip_boundary_dir = os.path.join(data_dir, "clip_boundaries")

feature_dir_file = set_feature_directions[dataset]

num_layers = set_num_layers[dataset]

feature_dirs_file = os.path.join(clip_boundary_dir, feature_dir_file)

boundary = "live_editing"

!python LARGE/get_clip_boundary.py --source_text "$source_text" --target_text "$target_text" --name $boundary --precomputed_dirs $feature_dirs_file --out_dir $clip_boundary_dir --cutoff_percentile $cutoff_percentile --model_layers $num_layers

boundary_path = os.path.join(clip_boundary_dir, boundary, f"boundary{boundary_ext}")

print(f"Done! The generated boundary can be found in: {boundary_path}")<jupyter_output><empty_output><jupyter_text>Upload your own boundary:The boundary should describe a direction in one of the latent spaces of StyleGAN2. We currently support boundaries in a numpy (npy) and pickle format, and in either W, W+ or S space.<jupyter_code>dataset = 'afhq_cat' #@param['celebA', 'afhq_cat']

boundary_path = "" #@param {type: "string"}

boundary_type = "w" #@param['w', 'w+', 's']
latent_src  = "psp" if boundary_type == "w+" else "e4e"
latent_type = "latents"<jupyter_output><empty_output><jupyter_text>2) Sort!Pick a seed and run the block in order to sort a random subset of CelebA according to your chosen semantic boundary.<jupyter_code>#@title Sort with the chosen boundary:

seed = 2 #@param {type: "integer"}
high_variance_samples = True #@param {type:"boolean"}

latents = cat_latents_dir if dataset == "afhq_cat" else os.path.join(latents_dir, latent_src, dataset, latent_type)
images  = cat_dir if dataset == "afhq_cat" else celeba_dir

output_dir = os.path.join("sorted", boundary)

if high_variance_samples:
  weighted_edges = True
  num_samples = 100
else:
  weighted_edges = False
  num_samples = 10

os.makedirs(output_dir, exist_ok=True)

num_layers = set_num_layers[dataset]

to_wp_flag = "--boundary_to_wp" if boundary_type == 'w' else ''
!python LARGE/sort.py --latent_path $latents --image_path $images --boundary_path $boundary_path --num_samples $num_samples --out_dir $output_dir --seed $seed --num_for_plot 10 $to_wp_flag --model_layers $num_layers --weighted_edges $weighted_edges
Image.open(os.path.join(output_dir, "all_sorted.jpg"))<jupyter_output><empty_output><jupyter_text>Regression For regression - first train a model using labeled CelebA data, then run inference on images from the CelebA test split or an uploaded image of your own.  Train a regression modelWe currently support pose and age models, using the labels from the paper<jupyter_code>model_type = 'pose' #@param['pose', 'age']

latents_path = os.path.join(latents_dir, "e4e", "celebA", "latents")

annotation_file = {"pose": "celeba_yaw.csv", "age": "celeba_age.csv"}[model_type]
boundary_path = {"pose": "pose_boundary", "age": "clip_boundaries/old_to_young_WP"}[model_type]
attribute = {"pose": "yaw", "age": "Young"}[model_type]

w_boundary_flags = f"--boundary_to_wp --layer_weights_path {data_dir}/{boundary_path}/layer_weights.npy" if model_type == "pose" else ""
distance_type_flags = "" if model_type == "pose" else "--distance_type euclidean"

!python LARGE/train_regression_model.py --data-path $latents_path --annotations-file $data_dir/labels/$annotation_file --boundary_path $data_dir/$boundary_path $w_boundary_flags --attribute $attribute --output-dir $data_dir/regression_models/ --regularization L1 L2 --feature_sample_ratio 0.95 $distance_type_flags<jupyter_output><empty_output><jupyter_text>Invert an image into the latent space of the GAN, or choose a pre-inverted image:  Choose a pre-inverted image (skip if you want to invert your own)<jupyter_code>image_files = os.listdir(celeba_dir)
image_paths = [os.path.join(celeba_dir, file_name) for file_name in image_files]

latent_type = "w+" #@param['w+', 's']
latent_dir_name = "s_latents" if latent_type == 's' else "latents"

latents = [os.path.join(latents_dir, "e4e", "celebA", latent_dir_name, file_name.split(".")[0] + ".pickle") for file_name in image_files]

image_id = 731 #@param {type:"slider", min:1, max:4500, step:1}

chosen_latent = latents[image_id - 1]
chosen_img = image_paths[image_id - 1]

img = Image.open(chosen_img)
img.resize((256, 256))<jupyter_output><empty_output><jupyter_text>Invert an image (skip if you picked a pre-inverted image)  Setup e4e for inversion. Skip if you want to use a pre-inverted image.<jupyter_code>#@title e4e setup
e4e_dir = os.path.join("/content", "encoder4editing")
sys.path.append(e4e_dir)

from models.psp import pSp
from gdown import download as drive_download

import torch
from torchvision import transforms
from argparse import Namespace

e4e_model_path = os.path.join(e4e_dir, "e4e_ffhq_encode.pt")
if not os.path.isfile(e4e_model_path):
    drive_download("https://drive.google.com/uc?id=1O8OLrVNOItOJoNGMyQ8G8YRTeTYEfs0P", e4e_model_path, quiet=False)
experiment_type = 'ffhq_encode'

# os.chdir('/content/encoder4editing')

EXPERIMENT_ARGS = {
        "model_path": e4e_model_path
    }
EXPERIMENT_ARGS['transform'] = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])
resize_dims = (256, 256)

model_path = EXPERIMENT_ARGS['model_path']
ckpt = torch.load(model_path, map_location='cpu')
opts = ckpt['opts']

# update the training options
opts['checkpoint_path'] = model_path
opts = Namespace(**opts)
e4e_net = pSp(opts)
e4e_net.eval()
e4e_net.cuda()
print('Model successfully loaded!')<jupyter_output><empty_output><jupyter_text>Choose and align an image for inversion. Skip if you want to use a pre-inverted image.<jupyter_code>#@title Align image
image_path = "latent_regression/sample_image/sample_steph.jpg" #@param {type: "string"}
original_image = Image.open(image_path)
original_image = original_image.convert("RGB")
if experiment_type == "ffhq_encode" and 'shape_predictor_68_face_landmarks.dat' not in os.listdir():
    !wget http://dlib.net/files/shape_predictor_68_face_landmarks.dat.bz2
    !bzip2 -dk shape_predictor_68_face_landmarks.dat.bz2

def run_alignment(image_path):
  import dlib
  from utils.alignment import align_face
  predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
  aligned_image = align_face(filepath=image_path, predictor=predictor) 
  print("Aligned image has shape: {}".format(aligned_image.size))
  return aligned_image 

if experiment_type == "ffhq_encode":
  input_image = run_alignment(image_path)
else:
  input_image = original_image

input_image.resize(resize_dims)<jupyter_output><empty_output><jupyter_text>Invert the aligned image. Skip if you want to use a pre-inverted image.<jupyter_code>#@title Invert the image

from utils.common import tensor2im

img_transforms = EXPERIMENT_ARGS['transform']
transformed_image = img_transforms(input_image)

def display_alongside_source_image(result_image, source_image):
    res = np.concatenate([np.array(source_image.resize(resize_dims)),
                          np.array(result_image.resize(resize_dims))], axis=1)
    return Image.fromarray(res)

def run_on_batch(inputs, net):
    images, latents = net(inputs.to("cuda").float(), randomize_noise=False, return_latents=True)
    if experiment_type == 'cars_encode':
        images = images[:, :, 32:224, :]
    return images, latents

with torch.no_grad():
    images, latents = run_on_batch(transformed_image.unsqueeze(0), e4e_net)
    result_image, latent = images[0], latents[0]

print(latents.size())

chosen_latent = os.path.join("/content", "latent.npy")
chosen_img = image_path

np.save(chosen_latent, latents.detach().cpu().numpy())

# Display inversion:
display_alongside_source_image(tensor2im(result_image), input_image)<jupyter_output><empty_output><jupyter_text>Predict semantic property strength using the trained regression model<jupyter_code>num_training_images = 1000 #@param[2, 5, 10, 20, 1000]

model_file = f"model_{num_training_images}.pickle"

!python LARGE/infer.py --pretrained_model $data_dir/regression_models/$model_file --latent_path "$chosen_latent" --boundary_path $data_dir/$boundary_path/boundary.npy $w_boundary_flags $distance_type_flags<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.11960989763842991, 'fraction_numerical': 0.012573547191101797, 'mean_word_length': 5.621131270010673, 'pattern_counts': {'":': 26, '<': 45, '<?xml version=': 0, '>': 45, 'https://': 6, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '379052', 'license': 'MIT', 'n_tokens_mistral': 4582, 'n_tokens_neox': 4356, 'n_words': 968, 'path': 'latent_regression.ipynb', 'repo_name': 'dumpmemory/LARGE', 'size': '12407'}