Duplicate from ucalyptus/PTI
Browse filesCo-authored-by: Sayantan Das <ucalyptus@users.noreply.huggingface.co>
This view is limited to 50 files because it contains too many changes.
See raw diff
- .gitattributes +27 -0
- .gitignore +5 -0
- README.md +38 -0
- app.py +26 -0
- configs/__init__.py +0 -0
- configs/__pycache__/__init__.cpython-36.pyc +0 -0
- configs/__pycache__/__init__.cpython-39.pyc +0 -0
- configs/__pycache__/global_config.cpython-36.pyc +0 -0
- configs/__pycache__/global_config.cpython-39.pyc +0 -0
- configs/__pycache__/hyperparameters.cpython-36.pyc +0 -0
- configs/__pycache__/hyperparameters.cpython-39.pyc +0 -0
- configs/__pycache__/paths_config.cpython-36.pyc +0 -0
- configs/__pycache__/paths_config.cpython-39.pyc +0 -0
- configs/evaluation_config.py +1 -0
- configs/global_config.py +12 -0
- configs/hyperparameters.py +28 -0
- configs/paths_config.py +31 -0
- criteria/__init__.py +0 -0
- criteria/__pycache__/__init__.cpython-36.pyc +0 -0
- criteria/__pycache__/__init__.cpython-39.pyc +0 -0
- criteria/__pycache__/l2_loss.cpython-36.pyc +0 -0
- criteria/__pycache__/l2_loss.cpython-39.pyc +0 -0
- criteria/__pycache__/localitly_regulizer.cpython-36.pyc +0 -0
- criteria/__pycache__/localitly_regulizer.cpython-39.pyc +0 -0
- criteria/l2_loss.py +8 -0
- criteria/localitly_regulizer.py +58 -0
- dnnlib/__init__.py +9 -0
- dnnlib/__pycache__/__init__.cpython-36.pyc +0 -0
- dnnlib/__pycache__/__init__.cpython-39.pyc +0 -0
- dnnlib/__pycache__/util.cpython-36.pyc +0 -0
- dnnlib/__pycache__/util.cpython-39.pyc +0 -0
- dnnlib/util.py +477 -0
- edit.py +83 -0
- editings/__pycache__/ganspace.cpython-39.pyc +0 -0
- editings/__pycache__/latent_editor.cpython-39.pyc +0 -0
- editings/ganspace.py +21 -0
- editings/ganspace_pca/ffhq_pca.pt +0 -0
- editings/interfacegan_directions/age.pt +0 -0
- editings/interfacegan_directions/rotation.pt +0 -0
- editings/interfacegan_directions/smile.pt +0 -0
- editings/latent_editor.py +23 -0
- evaluation/experiment_setting_creator.py +43 -0
- evaluation/qualitative_edit_comparison.py +156 -0
- makedirs.py +83 -0
- models/StyleCLIP/__init__.py +0 -0
- models/StyleCLIP/criteria/__init__.py +0 -0
- models/StyleCLIP/criteria/clip_loss.py +17 -0
- models/StyleCLIP/criteria/id_loss.py +39 -0
- models/StyleCLIP/global_directions/GUI.py +103 -0
- models/StyleCLIP/global_directions/GenerateImg.py +50 -0
.gitattributes
ADDED
@@ -0,0 +1,27 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
4 |
+
*.bin.* filter=lfs diff=lfs merge=lfs -text
|
5 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
11 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
12 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
13 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
14 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
15 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
16 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
17 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
18 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
19 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
20 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
21 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
22 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
23 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
24 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
25 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
26 |
+
*.zstandard filter=lfs diff=lfs merge=lfs -text
|
27 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
.gitignore
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
*.dat
|
2 |
+
*.pkl
|
3 |
+
*.zip
|
4 |
+
__pycache__/*
|
5 |
+
*.pyc
|
README.md
ADDED
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
title: PTI
|
3 |
+
emoji: 🦀
|
4 |
+
colorFrom: gray
|
5 |
+
colorTo: pink
|
6 |
+
sdk: gradio
|
7 |
+
app_file: app.py
|
8 |
+
pinned: false
|
9 |
+
duplicated_from: ucalyptus/PTI
|
10 |
+
---
|
11 |
+
|
12 |
+
# Configuration
|
13 |
+
|
14 |
+
`title`: _string_
|
15 |
+
Display title for the Space
|
16 |
+
|
17 |
+
`emoji`: _string_
|
18 |
+
Space emoji (emoji-only character allowed)
|
19 |
+
|
20 |
+
`colorFrom`: _string_
|
21 |
+
Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
|
22 |
+
|
23 |
+
`colorTo`: _string_
|
24 |
+
Color for Thumbnail gradient (red, yellow, green, blue, indigo, purple, pink, gray)
|
25 |
+
|
26 |
+
`sdk`: _string_
|
27 |
+
Can be either `gradio`, `streamlit`, or `static`
|
28 |
+
|
29 |
+
`sdk_version` : _string_
|
30 |
+
Only applicable for `streamlit` SDK.
|
31 |
+
See [doc](https://hf.co/docs/hub/spaces) for more info on supported versions.
|
32 |
+
|
33 |
+
`app_file`: _string_
|
34 |
+
Path to your main application file (which contains either `gradio` or `streamlit` Python code, or `static` html code).
|
35 |
+
Path is relative to the root of the repository.
|
36 |
+
|
37 |
+
`pinned`: _boolean_
|
38 |
+
Whether the Space stays on top of your list.
|
app.py
ADDED
@@ -0,0 +1,26 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
from PIL import Image
|
3 |
+
os.system("pip install gdown lpips gradio")
|
4 |
+
os.system("gdown --id 1HKmjg6iXsWr4aFPuU0gBXPGR83wqMzq7 -O align.dat")
|
5 |
+
os.system("wget https://nvlabs-fi-cdn.nvidia.com/stylegan2-ada-pytorch/pretrained/ffhq.pkl")
|
6 |
+
os.system("gdown https://github.com/ninja-build/ninja/releases/download/v1.10.2/ninja-linux.zip")
|
7 |
+
os.system("unzip ninja-linux.zip -d /usr/local/bin/")
|
8 |
+
os.system("sudo update-alternatives --install /usr/bin/ninja ninja /usr/local/bin/ninja 1 --force")
|
9 |
+
os.makedirs("embeddings/",exist_ok=True)
|
10 |
+
os.makedirs("checkpoints/",exist_ok=True)
|
11 |
+
os.makedirs("output/", exist_ok=True)
|
12 |
+
os.makedirs("images/",exist_ok=True)
|
13 |
+
import gradio as gr
|
14 |
+
|
15 |
+
|
16 |
+
|
17 |
+
def inference(img):
|
18 |
+
img.save("images/file.png")
|
19 |
+
os.system("python tune.py")
|
20 |
+
return Image.open("output/out.png")
|
21 |
+
|
22 |
+
title = "Pivotal Tuning for Latent-based editing of Real Images"
|
23 |
+
description = "Gradio Demo for Pivotal Tuning Inversion. To use it, simply upload your image. Ensure the image is 256 x 256 in size."
|
24 |
+
article = "<p style='text-align: center'><a href='https://github.com/danielroich/PTI' target='_blank'>Github Repo</a>"
|
25 |
+
gr.Interface(inference, [gr.inputs.Image(type="pil")], gr.outputs.Image(type="pil"),title=title,description=description,article=article,allow_flagging='never',allow_screenshot=False).launch(share=True)
|
26 |
+
os.system("rm output/out.png")
|
configs/__init__.py
ADDED
File without changes
|
configs/__pycache__/__init__.cpython-36.pyc
ADDED
Binary file (121 Bytes). View file
|
|
configs/__pycache__/__init__.cpython-39.pyc
ADDED
Binary file (129 Bytes). View file
|
|
configs/__pycache__/global_config.cpython-36.pyc
ADDED
Binary file (331 Bytes). View file
|
|
configs/__pycache__/global_config.cpython-39.pyc
ADDED
Binary file (339 Bytes). View file
|
|
configs/__pycache__/hyperparameters.cpython-36.pyc
ADDED
Binary file (680 Bytes). View file
|
|
configs/__pycache__/hyperparameters.cpython-39.pyc
ADDED
Binary file (693 Bytes). View file
|
|
configs/__pycache__/paths_config.cpython-36.pyc
ADDED
Binary file (1.06 kB). View file
|
|
configs/__pycache__/paths_config.cpython-39.pyc
ADDED
Binary file (1.04 kB). View file
|
|
configs/evaluation_config.py
ADDED
@@ -0,0 +1 @@
|
|
|
|
|
1 |
+
evaluated_methods = ['e4e', 'SG2', 'SG2Plus']
|
configs/global_config.py
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
## Device
|
2 |
+
cuda_visible_devices = '0'
|
3 |
+
device = 'cuda:0'
|
4 |
+
|
5 |
+
## Logs
|
6 |
+
training_step = 1
|
7 |
+
image_rec_result_log_snapshot = 100
|
8 |
+
pivotal_training_steps = 0
|
9 |
+
model_snapshot_interval = 400
|
10 |
+
|
11 |
+
## Run name to be updated during PTI
|
12 |
+
run_name = ''
|
configs/hyperparameters.py
ADDED
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
## Architechture
|
2 |
+
lpips_type = 'alex'
|
3 |
+
first_inv_type = 'w'
|
4 |
+
optim_type = 'adam'
|
5 |
+
|
6 |
+
## Locality regularization
|
7 |
+
latent_ball_num_of_samples = 1
|
8 |
+
locality_regularization_interval = 1
|
9 |
+
use_locality_regularization = False
|
10 |
+
regulizer_l2_lambda = 0.1
|
11 |
+
regulizer_lpips_lambda = 0.1
|
12 |
+
regulizer_alpha = 30
|
13 |
+
|
14 |
+
## Loss
|
15 |
+
pt_l2_lambda = 1
|
16 |
+
pt_lpips_lambda = 1
|
17 |
+
|
18 |
+
## Steps
|
19 |
+
LPIPS_value_threshold = 0.06
|
20 |
+
max_pti_steps = 350
|
21 |
+
first_inv_steps = 450
|
22 |
+
max_images_to_invert = 300
|
23 |
+
|
24 |
+
## Optimization
|
25 |
+
pti_learning_rate = 3e-4
|
26 |
+
first_inv_lr = 5e-3
|
27 |
+
train_batch_size = 1
|
28 |
+
use_last_w_pivots = False
|
configs/paths_config.py
ADDED
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
## Pretrained models paths
|
2 |
+
e4e = 'e4e_ffhq_encode.pt'
|
3 |
+
stylegan2_ada_ffhq = 'ffhq.pkl'
|
4 |
+
style_clip_pretrained_mappers = ''
|
5 |
+
ir_se50 = 'model_ir_se50.pth'
|
6 |
+
dlib = 'align.dat'
|
7 |
+
|
8 |
+
## Dirs for output files
|
9 |
+
checkpoints_dir = './checkpoints'
|
10 |
+
embedding_base_dir = './embeddings'
|
11 |
+
styleclip_output_dir = './StyleCLIP_results'
|
12 |
+
experiments_output_dir = './output'
|
13 |
+
|
14 |
+
## Input info
|
15 |
+
### Input dir, where the images reside
|
16 |
+
input_data_path = ''
|
17 |
+
### Inversion identifier, used to keeping track of the inversion results. Both the latent code and the generator
|
18 |
+
input_data_id = 'rocky'
|
19 |
+
|
20 |
+
## Keywords
|
21 |
+
pti_results_keyword = 'PTI'
|
22 |
+
e4e_results_keyword = 'e4e'
|
23 |
+
sg2_results_keyword = 'SG2'
|
24 |
+
sg2_plus_results_keyword = 'SG2_plus'
|
25 |
+
multi_id_model_type = 'multi_id'
|
26 |
+
|
27 |
+
## Edit directions
|
28 |
+
interfacegan_age = 'editings/interfacegan_directions/age.pt'
|
29 |
+
interfacegan_smile = 'editings/interfacegan_directions/smile.pt'
|
30 |
+
interfacegan_rotation = 'editings/interfacegan_directions/rotation.pt'
|
31 |
+
ffhq_pca = 'editings/ganspace_pca/ffhq_pca.pt'
|
criteria/__init__.py
ADDED
File without changes
|
criteria/__pycache__/__init__.cpython-36.pyc
ADDED
Binary file (122 Bytes). View file
|
|
criteria/__pycache__/__init__.cpython-39.pyc
ADDED
Binary file (130 Bytes). View file
|
|
criteria/__pycache__/l2_loss.cpython-36.pyc
ADDED
Binary file (346 Bytes). View file
|
|
criteria/__pycache__/l2_loss.cpython-39.pyc
ADDED
Binary file (358 Bytes). View file
|
|
criteria/__pycache__/localitly_regulizer.cpython-36.pyc
ADDED
Binary file (2.86 kB). View file
|
|
criteria/__pycache__/localitly_regulizer.cpython-39.pyc
ADDED
Binary file (2.9 kB). View file
|
|
criteria/l2_loss.py
ADDED
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
l2_criterion = torch.nn.MSELoss(reduction='mean')
|
4 |
+
|
5 |
+
|
6 |
+
def l2_loss(real_images, generated_images):
|
7 |
+
loss = l2_criterion(real_images, generated_images)
|
8 |
+
return loss
|
criteria/localitly_regulizer.py
ADDED
@@ -0,0 +1,58 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
import numpy as np
|
3 |
+
from criteria import l2_loss
|
4 |
+
from configs import hyperparameters
|
5 |
+
from configs import global_config
|
6 |
+
|
7 |
+
|
8 |
+
class Space_Regulizer:
|
9 |
+
def __init__(self, original_G, lpips_net):
|
10 |
+
self.original_G = original_G
|
11 |
+
self.morphing_regulizer_alpha = hyperparameters.regulizer_alpha
|
12 |
+
self.lpips_loss = lpips_net
|
13 |
+
|
14 |
+
def get_morphed_w_code(self, new_w_code, fixed_w):
|
15 |
+
interpolation_direction = new_w_code - fixed_w
|
16 |
+
interpolation_direction_norm = torch.norm(interpolation_direction, p=2)
|
17 |
+
direction_to_move = hyperparameters.regulizer_alpha * interpolation_direction / interpolation_direction_norm
|
18 |
+
result_w = fixed_w + direction_to_move
|
19 |
+
self.morphing_regulizer_alpha * fixed_w + (1 - self.morphing_regulizer_alpha) * new_w_code
|
20 |
+
|
21 |
+
return result_w
|
22 |
+
|
23 |
+
def get_image_from_ws(self, w_codes, G):
|
24 |
+
return torch.cat([G.synthesis(w_code, noise_mode='none', force_fp32=True) for w_code in w_codes])
|
25 |
+
|
26 |
+
def ball_holder_loss_lazy(self, new_G, num_of_sampled_latents, w_batch, use_wandb=False):
|
27 |
+
loss = 0.0
|
28 |
+
|
29 |
+
z_samples = np.random.randn(num_of_sampled_latents, self.original_G.z_dim)
|
30 |
+
w_samples = self.original_G.mapping(torch.from_numpy(z_samples).to(global_config.device), None,
|
31 |
+
truncation_psi=0.5)
|
32 |
+
territory_indicator_ws = [self.get_morphed_w_code(w_code.unsqueeze(0), w_batch) for w_code in w_samples]
|
33 |
+
|
34 |
+
for w_code in territory_indicator_ws:
|
35 |
+
new_img = new_G.synthesis(w_code, noise_mode='none', force_fp32=True)
|
36 |
+
with torch.no_grad():
|
37 |
+
old_img = self.original_G.synthesis(w_code, noise_mode='none', force_fp32=True)
|
38 |
+
|
39 |
+
if hyperparameters.regulizer_l2_lambda > 0:
|
40 |
+
l2_loss_val = l2_loss.l2_loss(old_img, new_img)
|
41 |
+
if use_wandb:
|
42 |
+
wandb.log({f'space_regulizer_l2_loss_val': l2_loss_val.detach().cpu()},
|
43 |
+
step=global_config.training_step)
|
44 |
+
loss += l2_loss_val * hyperparameters.regulizer_l2_lambda
|
45 |
+
|
46 |
+
if hyperparameters.regulizer_lpips_lambda > 0:
|
47 |
+
loss_lpips = self.lpips_loss(old_img, new_img)
|
48 |
+
loss_lpips = torch.mean(torch.squeeze(loss_lpips))
|
49 |
+
if use_wandb:
|
50 |
+
wandb.log({f'space_regulizer_lpips_loss_val': loss_lpips.detach().cpu()},
|
51 |
+
step=global_config.training_step)
|
52 |
+
loss += loss_lpips * hyperparameters.regulizer_lpips_lambda
|
53 |
+
|
54 |
+
return loss / len(territory_indicator_ws)
|
55 |
+
|
56 |
+
def space_regulizer_loss(self, new_G, w_batch, use_wandb):
|
57 |
+
ret_val = self.ball_holder_loss_lazy(new_G, hyperparameters.latent_ball_num_of_samples, w_batch, use_wandb)
|
58 |
+
return ret_val
|
dnnlib/__init__.py
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
2 |
+
#
|
3 |
+
# NVIDIA CORPORATION and its licensors retain all intellectual property
|
4 |
+
# and proprietary rights in and to this software, related documentation
|
5 |
+
# and any modifications thereto. Any use, reproduction, disclosure or
|
6 |
+
# distribution of this software and related documentation without an express
|
7 |
+
# license agreement from NVIDIA CORPORATION is strictly prohibited.
|
8 |
+
|
9 |
+
from .util import EasyDict, make_cache_dir_path
|
dnnlib/__pycache__/__init__.cpython-36.pyc
ADDED
Binary file (187 Bytes). View file
|
|
dnnlib/__pycache__/__init__.cpython-39.pyc
ADDED
Binary file (195 Bytes). View file
|
|
dnnlib/__pycache__/util.cpython-36.pyc
ADDED
Binary file (13.6 kB). View file
|
|
dnnlib/__pycache__/util.cpython-39.pyc
ADDED
Binary file (13.7 kB). View file
|
|
dnnlib/util.py
ADDED
@@ -0,0 +1,477 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved.
|
2 |
+
#
|
3 |
+
# NVIDIA CORPORATION and its licensors retain all intellectual property
|
4 |
+
# and proprietary rights in and to this software, related documentation
|
5 |
+
# and any modifications thereto. Any use, reproduction, disclosure or
|
6 |
+
# distribution of this software and related documentation without an express
|
7 |
+
# license agreement from NVIDIA CORPORATION is strictly prohibited.
|
8 |
+
|
9 |
+
"""Miscellaneous utility classes and functions."""
|
10 |
+
|
11 |
+
import ctypes
|
12 |
+
import fnmatch
|
13 |
+
import importlib
|
14 |
+
import inspect
|
15 |
+
import numpy as np
|
16 |
+
import os
|
17 |
+
import shutil
|
18 |
+
import sys
|
19 |
+
import types
|
20 |
+
import io
|
21 |
+
import pickle
|
22 |
+
import re
|
23 |
+
import requests
|
24 |
+
import html
|
25 |
+
import hashlib
|
26 |
+
import glob
|
27 |
+
import tempfile
|
28 |
+
import urllib
|
29 |
+
import urllib.request
|
30 |
+
import uuid
|
31 |
+
|
32 |
+
from distutils.util import strtobool
|
33 |
+
from typing import Any, List, Tuple, Union
|
34 |
+
|
35 |
+
|
36 |
+
# Util classes
|
37 |
+
# ------------------------------------------------------------------------------------------
|
38 |
+
|
39 |
+
|
40 |
+
class EasyDict(dict):
|
41 |
+
"""Convenience class that behaves like a dict but allows access with the attribute syntax."""
|
42 |
+
|
43 |
+
def __getattr__(self, name: str) -> Any:
|
44 |
+
try:
|
45 |
+
return self[name]
|
46 |
+
except KeyError:
|
47 |
+
raise AttributeError(name)
|
48 |
+
|
49 |
+
def __setattr__(self, name: str, value: Any) -> None:
|
50 |
+
self[name] = value
|
51 |
+
|
52 |
+
def __delattr__(self, name: str) -> None:
|
53 |
+
del self[name]
|
54 |
+
|
55 |
+
|
56 |
+
class Logger(object):
|
57 |
+
"""Redirect stderr to stdout, optionally print stdout to a file, and optionally force flushing on both stdout and the file."""
|
58 |
+
|
59 |
+
def __init__(self, file_name: str = None, file_mode: str = "w", should_flush: bool = True):
|
60 |
+
self.file = None
|
61 |
+
|
62 |
+
if file_name is not None:
|
63 |
+
self.file = open(file_name, file_mode)
|
64 |
+
|
65 |
+
self.should_flush = should_flush
|
66 |
+
self.stdout = sys.stdout
|
67 |
+
self.stderr = sys.stderr
|
68 |
+
|
69 |
+
sys.stdout = self
|
70 |
+
sys.stderr = self
|
71 |
+
|
72 |
+
def __enter__(self) -> "Logger":
|
73 |
+
return self
|
74 |
+
|
75 |
+
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
|
76 |
+
self.close()
|
77 |
+
|
78 |
+
def write(self, text: Union[str, bytes]) -> None:
|
79 |
+
"""Write text to stdout (and a file) and optionally flush."""
|
80 |
+
if isinstance(text, bytes):
|
81 |
+
text = text.decode()
|
82 |
+
if len(text) == 0: # workaround for a bug in VSCode debugger: sys.stdout.write(''); sys.stdout.flush() => crash
|
83 |
+
return
|
84 |
+
|
85 |
+
if self.file is not None:
|
86 |
+
self.file.write(text)
|
87 |
+
|
88 |
+
self.stdout.write(text)
|
89 |
+
|
90 |
+
if self.should_flush:
|
91 |
+
self.flush()
|
92 |
+
|
93 |
+
def flush(self) -> None:
|
94 |
+
"""Flush written text to both stdout and a file, if open."""
|
95 |
+
if self.file is not None:
|
96 |
+
self.file.flush()
|
97 |
+
|
98 |
+
self.stdout.flush()
|
99 |
+
|
100 |
+
def close(self) -> None:
|
101 |
+
"""Flush, close possible files, and remove stdout/stderr mirroring."""
|
102 |
+
self.flush()
|
103 |
+
|
104 |
+
# if using multiple loggers, prevent closing in wrong order
|
105 |
+
if sys.stdout is self:
|
106 |
+
sys.stdout = self.stdout
|
107 |
+
if sys.stderr is self:
|
108 |
+
sys.stderr = self.stderr
|
109 |
+
|
110 |
+
if self.file is not None:
|
111 |
+
self.file.close()
|
112 |
+
self.file = None
|
113 |
+
|
114 |
+
|
115 |
+
# Cache directories
|
116 |
+
# ------------------------------------------------------------------------------------------
|
117 |
+
|
118 |
+
_dnnlib_cache_dir = None
|
119 |
+
|
120 |
+
def set_cache_dir(path: str) -> None:
|
121 |
+
global _dnnlib_cache_dir
|
122 |
+
_dnnlib_cache_dir = path
|
123 |
+
|
124 |
+
def make_cache_dir_path(*paths: str) -> str:
|
125 |
+
if _dnnlib_cache_dir is not None:
|
126 |
+
return os.path.join(_dnnlib_cache_dir, *paths)
|
127 |
+
if 'DNNLIB_CACHE_DIR' in os.environ:
|
128 |
+
return os.path.join(os.environ['DNNLIB_CACHE_DIR'], *paths)
|
129 |
+
if 'HOME' in os.environ:
|
130 |
+
return os.path.join(os.environ['HOME'], '.cache', 'dnnlib', *paths)
|
131 |
+
if 'USERPROFILE' in os.environ:
|
132 |
+
return os.path.join(os.environ['USERPROFILE'], '.cache', 'dnnlib', *paths)
|
133 |
+
return os.path.join(tempfile.gettempdir(), '.cache', 'dnnlib', *paths)
|
134 |
+
|
135 |
+
# Small util functions
|
136 |
+
# ------------------------------------------------------------------------------------------
|
137 |
+
|
138 |
+
|
139 |
+
def format_time(seconds: Union[int, float]) -> str:
|
140 |
+
"""Convert the seconds to human readable string with days, hours, minutes and seconds."""
|
141 |
+
s = int(np.rint(seconds))
|
142 |
+
|
143 |
+
if s < 60:
|
144 |
+
return "{0}s".format(s)
|
145 |
+
elif s < 60 * 60:
|
146 |
+
return "{0}m {1:02}s".format(s // 60, s % 60)
|
147 |
+
elif s < 24 * 60 * 60:
|
148 |
+
return "{0}h {1:02}m {2:02}s".format(s // (60 * 60), (s // 60) % 60, s % 60)
|
149 |
+
else:
|
150 |
+
return "{0}d {1:02}h {2:02}m".format(s // (24 * 60 * 60), (s // (60 * 60)) % 24, (s // 60) % 60)
|
151 |
+
|
152 |
+
|
153 |
+
def ask_yes_no(question: str) -> bool:
|
154 |
+
"""Ask the user the question until the user inputs a valid answer."""
|
155 |
+
while True:
|
156 |
+
try:
|
157 |
+
print("{0} [y/n]".format(question))
|
158 |
+
return strtobool(input().lower())
|
159 |
+
except ValueError:
|
160 |
+
pass
|
161 |
+
|
162 |
+
|
163 |
+
def tuple_product(t: Tuple) -> Any:
|
164 |
+
"""Calculate the product of the tuple elements."""
|
165 |
+
result = 1
|
166 |
+
|
167 |
+
for v in t:
|
168 |
+
result *= v
|
169 |
+
|
170 |
+
return result
|
171 |
+
|
172 |
+
|
173 |
+
_str_to_ctype = {
|
174 |
+
"uint8": ctypes.c_ubyte,
|
175 |
+
"uint16": ctypes.c_uint16,
|
176 |
+
"uint32": ctypes.c_uint32,
|
177 |
+
"uint64": ctypes.c_uint64,
|
178 |
+
"int8": ctypes.c_byte,
|
179 |
+
"int16": ctypes.c_int16,
|
180 |
+
"int32": ctypes.c_int32,
|
181 |
+
"int64": ctypes.c_int64,
|
182 |
+
"float32": ctypes.c_float,
|
183 |
+
"float64": ctypes.c_double
|
184 |
+
}
|
185 |
+
|
186 |
+
|
187 |
+
def get_dtype_and_ctype(type_obj: Any) -> Tuple[np.dtype, Any]:
|
188 |
+
"""Given a type name string (or an object having a __name__ attribute), return matching Numpy and ctypes types that have the same size in bytes."""
|
189 |
+
type_str = None
|
190 |
+
|
191 |
+
if isinstance(type_obj, str):
|
192 |
+
type_str = type_obj
|
193 |
+
elif hasattr(type_obj, "__name__"):
|
194 |
+
type_str = type_obj.__name__
|
195 |
+
elif hasattr(type_obj, "name"):
|
196 |
+
type_str = type_obj.name
|
197 |
+
else:
|
198 |
+
raise RuntimeError("Cannot infer type name from input")
|
199 |
+
|
200 |
+
assert type_str in _str_to_ctype.keys()
|
201 |
+
|
202 |
+
my_dtype = np.dtype(type_str)
|
203 |
+
my_ctype = _str_to_ctype[type_str]
|
204 |
+
|
205 |
+
assert my_dtype.itemsize == ctypes.sizeof(my_ctype)
|
206 |
+
|
207 |
+
return my_dtype, my_ctype
|
208 |
+
|
209 |
+
|
210 |
+
def is_pickleable(obj: Any) -> bool:
|
211 |
+
try:
|
212 |
+
with io.BytesIO() as stream:
|
213 |
+
pickle.dump(obj, stream)
|
214 |
+
return True
|
215 |
+
except:
|
216 |
+
return False
|
217 |
+
|
218 |
+
|
219 |
+
# Functionality to import modules/objects by name, and call functions by name
|
220 |
+
# ------------------------------------------------------------------------------------------
|
221 |
+
|
222 |
+
def get_module_from_obj_name(obj_name: str) -> Tuple[types.ModuleType, str]:
|
223 |
+
"""Searches for the underlying module behind the name to some python object.
|
224 |
+
Returns the module and the object name (original name with module part removed)."""
|
225 |
+
|
226 |
+
# allow convenience shorthands, substitute them by full names
|
227 |
+
obj_name = re.sub("^np.", "numpy.", obj_name)
|
228 |
+
obj_name = re.sub("^tf.", "tensorflow.", obj_name)
|
229 |
+
|
230 |
+
# list alternatives for (module_name, local_obj_name)
|
231 |
+
parts = obj_name.split(".")
|
232 |
+
name_pairs = [(".".join(parts[:i]), ".".join(parts[i:])) for i in range(len(parts), 0, -1)]
|
233 |
+
|
234 |
+
# try each alternative in turn
|
235 |
+
for module_name, local_obj_name in name_pairs:
|
236 |
+
try:
|
237 |
+
module = importlib.import_module(module_name) # may raise ImportError
|
238 |
+
get_obj_from_module(module, local_obj_name) # may raise AttributeError
|
239 |
+
return module, local_obj_name
|
240 |
+
except:
|
241 |
+
pass
|
242 |
+
|
243 |
+
# maybe some of the modules themselves contain errors?
|
244 |
+
for module_name, _local_obj_name in name_pairs:
|
245 |
+
try:
|
246 |
+
importlib.import_module(module_name) # may raise ImportError
|
247 |
+
except ImportError:
|
248 |
+
if not str(sys.exc_info()[1]).startswith("No module named '" + module_name + "'"):
|
249 |
+
raise
|
250 |
+
|
251 |
+
# maybe the requested attribute is missing?
|
252 |
+
for module_name, local_obj_name in name_pairs:
|
253 |
+
try:
|
254 |
+
module = importlib.import_module(module_name) # may raise ImportError
|
255 |
+
get_obj_from_module(module, local_obj_name) # may raise AttributeError
|
256 |
+
except ImportError:
|
257 |
+
pass
|
258 |
+
|
259 |
+
# we are out of luck, but we have no idea why
|
260 |
+
raise ImportError(obj_name)
|
261 |
+
|
262 |
+
|
263 |
+
def get_obj_from_module(module: types.ModuleType, obj_name: str) -> Any:
|
264 |
+
"""Traverses the object name and returns the last (rightmost) python object."""
|
265 |
+
if obj_name == '':
|
266 |
+
return module
|
267 |
+
obj = module
|
268 |
+
for part in obj_name.split("."):
|
269 |
+
obj = getattr(obj, part)
|
270 |
+
return obj
|
271 |
+
|
272 |
+
|
273 |
+
def get_obj_by_name(name: str) -> Any:
|
274 |
+
"""Finds the python object with the given name."""
|
275 |
+
module, obj_name = get_module_from_obj_name(name)
|
276 |
+
return get_obj_from_module(module, obj_name)
|
277 |
+
|
278 |
+
|
279 |
+
def call_func_by_name(*args, func_name: str = None, **kwargs) -> Any:
|
280 |
+
"""Finds the python object with the given name and calls it as a function."""
|
281 |
+
assert func_name is not None
|
282 |
+
func_obj = get_obj_by_name(func_name)
|
283 |
+
assert callable(func_obj)
|
284 |
+
return func_obj(*args, **kwargs)
|
285 |
+
|
286 |
+
|
287 |
+
def construct_class_by_name(*args, class_name: str = None, **kwargs) -> Any:
|
288 |
+
"""Finds the python class with the given name and constructs it with the given arguments."""
|
289 |
+
return call_func_by_name(*args, func_name=class_name, **kwargs)
|
290 |
+
|
291 |
+
|
292 |
+
def get_module_dir_by_obj_name(obj_name: str) -> str:
|
293 |
+
"""Get the directory path of the module containing the given object name."""
|
294 |
+
module, _ = get_module_from_obj_name(obj_name)
|
295 |
+
return os.path.dirname(inspect.getfile(module))
|
296 |
+
|
297 |
+
|
298 |
+
def is_top_level_function(obj: Any) -> bool:
|
299 |
+
"""Determine whether the given object is a top-level function, i.e., defined at module scope using 'def'."""
|
300 |
+
return callable(obj) and obj.__name__ in sys.modules[obj.__module__].__dict__
|
301 |
+
|
302 |
+
|
303 |
+
def get_top_level_function_name(obj: Any) -> str:
|
304 |
+
"""Return the fully-qualified name of a top-level function."""
|
305 |
+
assert is_top_level_function(obj)
|
306 |
+
module = obj.__module__
|
307 |
+
if module == '__main__':
|
308 |
+
module = os.path.splitext(os.path.basename(sys.modules[module].__file__))[0]
|
309 |
+
return module + "." + obj.__name__
|
310 |
+
|
311 |
+
|
312 |
+
# File system helpers
|
313 |
+
# ------------------------------------------------------------------------------------------
|
314 |
+
|
315 |
+
def list_dir_recursively_with_ignore(dir_path: str, ignores: List[str] = None, add_base_to_relative: bool = False) -> List[Tuple[str, str]]:
|
316 |
+
"""List all files recursively in a given directory while ignoring given file and directory names.
|
317 |
+
Returns list of tuples containing both absolute and relative paths."""
|
318 |
+
assert os.path.isdir(dir_path)
|
319 |
+
base_name = os.path.basename(os.path.normpath(dir_path))
|
320 |
+
|
321 |
+
if ignores is None:
|
322 |
+
ignores = []
|
323 |
+
|
324 |
+
result = []
|
325 |
+
|
326 |
+
for root, dirs, files in os.walk(dir_path, topdown=True):
|
327 |
+
for ignore_ in ignores:
|
328 |
+
dirs_to_remove = [d for d in dirs if fnmatch.fnmatch(d, ignore_)]
|
329 |
+
|
330 |
+
# dirs need to be edited in-place
|
331 |
+
for d in dirs_to_remove:
|
332 |
+
dirs.remove(d)
|
333 |
+
|
334 |
+
files = [f for f in files if not fnmatch.fnmatch(f, ignore_)]
|
335 |
+
|
336 |
+
absolute_paths = [os.path.join(root, f) for f in files]
|
337 |
+
relative_paths = [os.path.relpath(p, dir_path) for p in absolute_paths]
|
338 |
+
|
339 |
+
if add_base_to_relative:
|
340 |
+
relative_paths = [os.path.join(base_name, p) for p in relative_paths]
|
341 |
+
|
342 |
+
assert len(absolute_paths) == len(relative_paths)
|
343 |
+
result += zip(absolute_paths, relative_paths)
|
344 |
+
|
345 |
+
return result
|
346 |
+
|
347 |
+
|
348 |
+
def copy_files_and_create_dirs(files: List[Tuple[str, str]]) -> None:
|
349 |
+
"""Takes in a list of tuples of (src, dst) paths and copies files.
|
350 |
+
Will create all necessary directories."""
|
351 |
+
for file in files:
|
352 |
+
target_dir_name = os.path.dirname(file[1])
|
353 |
+
|
354 |
+
# will create all intermediate-level directories
|
355 |
+
if not os.path.exists(target_dir_name):
|
356 |
+
os.makedirs(target_dir_name)
|
357 |
+
|
358 |
+
shutil.copyfile(file[0], file[1])
|
359 |
+
|
360 |
+
|
361 |
+
# URL helpers
|
362 |
+
# ------------------------------------------------------------------------------------------
|
363 |
+
|
364 |
+
def is_url(obj: Any, allow_file_urls: bool = False) -> bool:
|
365 |
+
"""Determine whether the given object is a valid URL string."""
|
366 |
+
if not isinstance(obj, str) or not "://" in obj:
|
367 |
+
return False
|
368 |
+
if allow_file_urls and obj.startswith('file://'):
|
369 |
+
return True
|
370 |
+
try:
|
371 |
+
res = requests.compat.urlparse(obj)
|
372 |
+
if not res.scheme or not res.netloc or not "." in res.netloc:
|
373 |
+
return False
|
374 |
+
res = requests.compat.urlparse(requests.compat.urljoin(obj, "/"))
|
375 |
+
if not res.scheme or not res.netloc or not "." in res.netloc:
|
376 |
+
return False
|
377 |
+
except:
|
378 |
+
return False
|
379 |
+
return True
|
380 |
+
|
381 |
+
|
382 |
+
def open_url(url: str, cache_dir: str = None, num_attempts: int = 10, verbose: bool = True, return_filename: bool = False, cache: bool = True) -> Any:
|
383 |
+
"""Download the given URL and return a binary-mode file object to access the data."""
|
384 |
+
assert num_attempts >= 1
|
385 |
+
assert not (return_filename and (not cache))
|
386 |
+
|
387 |
+
# Doesn't look like an URL scheme so interpret it as a local filename.
|
388 |
+
if not re.match('^[a-z]+://', url):
|
389 |
+
return url if return_filename else open(url, "rb")
|
390 |
+
|
391 |
+
# Handle file URLs. This code handles unusual file:// patterns that
|
392 |
+
# arise on Windows:
|
393 |
+
#
|
394 |
+
# file:///c:/foo.txt
|
395 |
+
#
|
396 |
+
# which would translate to a local '/c:/foo.txt' filename that's
|
397 |
+
# invalid. Drop the forward slash for such pathnames.
|
398 |
+
#
|
399 |
+
# If you touch this code path, you should test it on both Linux and
|
400 |
+
# Windows.
|
401 |
+
#
|
402 |
+
# Some internet resources suggest using urllib.request.url2pathname() but
|
403 |
+
# but that converts forward slashes to backslashes and this causes
|
404 |
+
# its own set of problems.
|
405 |
+
if url.startswith('file://'):
|
406 |
+
filename = urllib.parse.urlparse(url).path
|
407 |
+
if re.match(r'^/[a-zA-Z]:', filename):
|
408 |
+
filename = filename[1:]
|
409 |
+
return filename if return_filename else open(filename, "rb")
|
410 |
+
|
411 |
+
assert is_url(url)
|
412 |
+
|
413 |
+
# Lookup from cache.
|
414 |
+
if cache_dir is None:
|
415 |
+
cache_dir = make_cache_dir_path('downloads')
|
416 |
+
|
417 |
+
url_md5 = hashlib.md5(url.encode("utf-8")).hexdigest()
|
418 |
+
if cache:
|
419 |
+
cache_files = glob.glob(os.path.join(cache_dir, url_md5 + "_*"))
|
420 |
+
if len(cache_files) == 1:
|
421 |
+
filename = cache_files[0]
|
422 |
+
return filename if return_filename else open(filename, "rb")
|
423 |
+
|
424 |
+
# Download.
|
425 |
+
url_name = None
|
426 |
+
url_data = None
|
427 |
+
with requests.Session() as session:
|
428 |
+
if verbose:
|
429 |
+
print("Downloading %s ..." % url, end="", flush=True)
|
430 |
+
for attempts_left in reversed(range(num_attempts)):
|
431 |
+
try:
|
432 |
+
with session.get(url) as res:
|
433 |
+
res.raise_for_status()
|
434 |
+
if len(res.content) == 0:
|
435 |
+
raise IOError("No data received")
|
436 |
+
|
437 |
+
if len(res.content) < 8192:
|
438 |
+
content_str = res.content.decode("utf-8")
|
439 |
+
if "download_warning" in res.headers.get("Set-Cookie", ""):
|
440 |
+
links = [html.unescape(link) for link in content_str.split('"') if "export=download" in link]
|
441 |
+
if len(links) == 1:
|
442 |
+
url = requests.compat.urljoin(url, links[0])
|
443 |
+
raise IOError("Google Drive virus checker nag")
|
444 |
+
if "Google Drive - Quota exceeded" in content_str:
|
445 |
+
raise IOError("Google Drive download quota exceeded -- please try again later")
|
446 |
+
|
447 |
+
match = re.search(r'filename="([^"]*)"', res.headers.get("Content-Disposition", ""))
|
448 |
+
url_name = match[1] if match else url
|
449 |
+
url_data = res.content
|
450 |
+
if verbose:
|
451 |
+
print(" done")
|
452 |
+
break
|
453 |
+
except KeyboardInterrupt:
|
454 |
+
raise
|
455 |
+
except:
|
456 |
+
if not attempts_left:
|
457 |
+
if verbose:
|
458 |
+
print(" failed")
|
459 |
+
raise
|
460 |
+
if verbose:
|
461 |
+
print(".", end="", flush=True)
|
462 |
+
|
463 |
+
# Save to cache.
|
464 |
+
if cache:
|
465 |
+
safe_name = re.sub(r"[^0-9a-zA-Z-._]", "_", url_name)
|
466 |
+
cache_file = os.path.join(cache_dir, url_md5 + "_" + safe_name)
|
467 |
+
temp_file = os.path.join(cache_dir, "tmp_" + uuid.uuid4().hex + "_" + url_md5 + "_" + safe_name)
|
468 |
+
os.makedirs(cache_dir, exist_ok=True)
|
469 |
+
with open(temp_file, "wb") as f:
|
470 |
+
f.write(url_data)
|
471 |
+
os.replace(temp_file, cache_file) # atomic
|
472 |
+
if return_filename:
|
473 |
+
return cache_file
|
474 |
+
|
475 |
+
# Return data as file object.
|
476 |
+
assert not return_filename
|
477 |
+
return io.BytesIO(url_data)
|
edit.py
ADDED
@@ -0,0 +1,83 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import click
|
2 |
+
import os
|
3 |
+
import sys
|
4 |
+
import pickle
|
5 |
+
import numpy as np
|
6 |
+
from PIL import Image
|
7 |
+
import glob
|
8 |
+
import torch
|
9 |
+
from configs import paths_config, hyperparameters, global_config
|
10 |
+
from IPython.display import display
|
11 |
+
import matplotlib.pyplot as plt
|
12 |
+
from scripts.latent_editor_wrapper import LatentEditorWrapper
|
13 |
+
|
14 |
+
|
15 |
+
image_dir_name = '/home/sayantan/processed_images'
|
16 |
+
use_multi_id_training = False
|
17 |
+
global_config.device = 'cuda'
|
18 |
+
paths_config.e4e = '/home/sayantan/PTI/pretrained_models/e4e_ffhq_encode.pt'
|
19 |
+
paths_config.input_data_id = image_dir_name
|
20 |
+
paths_config.input_data_path = f'{image_dir_name}'
|
21 |
+
paths_config.stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
|
22 |
+
paths_config.checkpoints_dir = '/home/sayantan/PTI/'
|
23 |
+
paths_config.style_clip_pretrained_mappers = '/home/sayantan/PTI/pretrained_models'
|
24 |
+
hyperparameters.use_locality_regularization = False
|
25 |
+
hyperparameters.lpips_type = 'squeeze'
|
26 |
+
|
27 |
+
model_id = "MYJJDFVGATAT"
|
28 |
+
|
29 |
+
|
30 |
+
|
31 |
+
def display_alongside_source_image(images):
|
32 |
+
res = np.concatenate([np.array(image) for image in images], axis=1)
|
33 |
+
return Image.fromarray(res)
|
34 |
+
|
35 |
+
def load_generators(model_id, image_name):
|
36 |
+
with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
|
37 |
+
old_G = pickle.load(f)['G_ema'].cuda()
|
38 |
+
|
39 |
+
with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new:
|
40 |
+
new_G = torch.load(f_new).cuda()
|
41 |
+
|
42 |
+
return old_G, new_G
|
43 |
+
|
44 |
+
def plot_syn_images(syn_images,text):
|
45 |
+
for img in syn_images:
|
46 |
+
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
|
47 |
+
plt.axis('off')
|
48 |
+
resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
|
49 |
+
display(resized_image)
|
50 |
+
#wandb.log({text: [wandb.Image(resized_image, caption="Label")]})
|
51 |
+
del img
|
52 |
+
del resized_image
|
53 |
+
torch.cuda.empty_cache()
|
54 |
+
|
55 |
+
def syn_images_wandb(img):
|
56 |
+
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
|
57 |
+
plt.axis('off')
|
58 |
+
resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
|
59 |
+
return resized_image
|
60 |
+
|
61 |
+
|
62 |
+
def edit(image_name):
|
63 |
+
generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
|
64 |
+
old_G, new_G = load_generators(model_id, generator_type)
|
65 |
+
w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
|
66 |
+
|
67 |
+
embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
|
68 |
+
w_pivot = torch.load(f'{embedding_dir}/0.pt')
|
69 |
+
|
70 |
+
old_image = old_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
|
71 |
+
new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
|
72 |
+
|
73 |
+
latent_editor = LatentEditorWrapper()
|
74 |
+
latents_after_edit = latent_editor.get_single_interface_gan_edits(w_pivot, [i for i in range(-5,5)])
|
75 |
+
|
76 |
+
for direction, factor_and_edit in latents_after_edit.items():
|
77 |
+
for editkey in factor_and_edit.keys():
|
78 |
+
new_image = new_G.synthesis(factor_and_edit[editkey], noise_mode='const', force_fp32 = True)
|
79 |
+
image_pil = syn_images_wandb(new_image).save(f"/home/sayantan/PTI/{direction}/{editkey}/{image_name}.jpg")
|
80 |
+
|
81 |
+
if __name__ == '__main__':
|
82 |
+
for image_name in [f.split(".")[0].split("_")[2] for f in sorted(glob.glob("*.pt"))]:
|
83 |
+
edit(image_name)
|
editings/__pycache__/ganspace.cpython-39.pyc
ADDED
Binary file (878 Bytes). View file
|
|
editings/__pycache__/latent_editor.cpython-39.pyc
ADDED
Binary file (994 Bytes). View file
|
|
editings/ganspace.py
ADDED
@@ -0,0 +1,21 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
|
4 |
+
def edit(latents, pca, edit_directions):
|
5 |
+
edit_latents = []
|
6 |
+
for latent in latents:
|
7 |
+
for pca_idx, start, end, strength in edit_directions:
|
8 |
+
delta = get_delta(pca, latent, pca_idx, strength)
|
9 |
+
delta_padded = torch.zeros(latent.shape).to('cuda')
|
10 |
+
delta_padded[start:end] += delta.repeat(end - start, 1)
|
11 |
+
edit_latents.append(latent + delta_padded)
|
12 |
+
return torch.stack(edit_latents)
|
13 |
+
|
14 |
+
|
15 |
+
def get_delta(pca, latent, idx, strength):
|
16 |
+
w_centered = latent - pca['mean'].to('cuda')
|
17 |
+
lat_comp = pca['comp'].to('cuda')
|
18 |
+
lat_std = pca['std'].to('cuda')
|
19 |
+
w_coord = torch.sum(w_centered[0].reshape(-1)*lat_comp[idx].reshape(-1)) / lat_std[idx]
|
20 |
+
delta = (strength - w_coord)*lat_comp[idx]*lat_std[idx]
|
21 |
+
return delta
|
editings/ganspace_pca/ffhq_pca.pt
ADDED
Binary file (168 kB). View file
|
|
editings/interfacegan_directions/age.pt
ADDED
Binary file (2.81 kB). View file
|
|
editings/interfacegan_directions/rotation.pt
ADDED
Binary file (2.81 kB). View file
|
|
editings/interfacegan_directions/smile.pt
ADDED
Binary file (2.81 kB). View file
|
|
editings/latent_editor.py
ADDED
@@ -0,0 +1,23 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
|
3 |
+
from configs import paths_config
|
4 |
+
from editings import ganspace
|
5 |
+
from utils.data_utils import tensor2im
|
6 |
+
|
7 |
+
|
8 |
+
class LatentEditor(object):
|
9 |
+
|
10 |
+
def apply_ganspace(self, latent, ganspace_pca, edit_directions):
|
11 |
+
edit_latents = ganspace.edit(latent, ganspace_pca, edit_directions)
|
12 |
+
return edit_latents
|
13 |
+
|
14 |
+
def apply_interfacegan(self, latent, direction, factor=1, factor_range=None):
|
15 |
+
edit_latents = []
|
16 |
+
if factor_range is not None: # Apply a range of editing factors. for example, (-5, 5)
|
17 |
+
for f in range(*factor_range):
|
18 |
+
edit_latent = latent + f * direction
|
19 |
+
edit_latents.append(edit_latent)
|
20 |
+
edit_latents = torch.cat(edit_latents)
|
21 |
+
else:
|
22 |
+
edit_latents = latent + factor * direction
|
23 |
+
return edit_latents
|
evaluation/experiment_setting_creator.py
ADDED
@@ -0,0 +1,43 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import glob
|
2 |
+
import os
|
3 |
+
from configs import global_config, paths_config, hyperparameters
|
4 |
+
from scripts.latent_creators.sg2_plus_latent_creator import SG2PlusLatentCreator
|
5 |
+
from scripts.latent_creators.e4e_latent_creator import E4ELatentCreator
|
6 |
+
from scripts.run_pti import run_PTI
|
7 |
+
import pickle
|
8 |
+
import torch
|
9 |
+
from utils.models_utils import toogle_grad, load_old_G
|
10 |
+
|
11 |
+
|
12 |
+
class ExperimentRunner:
|
13 |
+
|
14 |
+
def __init__(self, run_id=''):
|
15 |
+
self.images_paths = glob.glob(f'{paths_config.input_data_path}/*')
|
16 |
+
self.target_paths = glob.glob(f'{paths_config.input_data_path}/*')
|
17 |
+
self.run_id = run_id
|
18 |
+
self.sampled_ws = None
|
19 |
+
|
20 |
+
self.old_G = load_old_G()
|
21 |
+
|
22 |
+
toogle_grad(self.old_G, False)
|
23 |
+
|
24 |
+
def run_experiment(self, run_pt, create_other_latents, use_multi_id_training, use_wandb=False):
|
25 |
+
if run_pt:
|
26 |
+
self.run_id = run_PTI(self.run_id, use_wandb=use_wandb, use_multi_id_training=use_multi_id_training)
|
27 |
+
if create_other_latents:
|
28 |
+
sg2_plus_latent_creator = SG2PlusLatentCreator(use_wandb=use_wandb)
|
29 |
+
sg2_plus_latent_creator.create_latents()
|
30 |
+
e4e_latent_creator = E4ELatentCreator(use_wandb=use_wandb)
|
31 |
+
e4e_latent_creator.create_latents()
|
32 |
+
|
33 |
+
torch.cuda.empty_cache()
|
34 |
+
|
35 |
+
return self.run_id
|
36 |
+
|
37 |
+
|
38 |
+
if __name__ == '__main__':
|
39 |
+
os.environ['CUDA_DEVICE_ORDER'] = 'PCI_BUS_ID'
|
40 |
+
os.environ['CUDA_VISIBLE_DEVICES'] = global_config.cuda_visible_devices
|
41 |
+
|
42 |
+
runner = ExperimentRunner()
|
43 |
+
runner.run_experiment(True, False, False)
|
evaluation/qualitative_edit_comparison.py
ADDED
@@ -0,0 +1,156 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
from random import choice
|
3 |
+
from string import ascii_uppercase
|
4 |
+
from PIL import Image
|
5 |
+
from tqdm import tqdm
|
6 |
+
from scripts.latent_editor_wrapper import LatentEditorWrapper
|
7 |
+
from evaluation.experiment_setting_creator import ExperimentRunner
|
8 |
+
import torch
|
9 |
+
from configs import paths_config, hyperparameters, evaluation_config
|
10 |
+
from utils.log_utils import save_concat_image, save_single_image
|
11 |
+
from utils.models_utils import load_tuned_G
|
12 |
+
|
13 |
+
|
14 |
+
class EditComparison:
|
15 |
+
|
16 |
+
def __init__(self, save_single_images, save_concatenated_images, run_id):
|
17 |
+
|
18 |
+
self.run_id = run_id
|
19 |
+
self.experiment_creator = ExperimentRunner(run_id)
|
20 |
+
self.save_single_images = save_single_images
|
21 |
+
self.save_concatenated_images = save_concatenated_images
|
22 |
+
self.latent_editor = LatentEditorWrapper()
|
23 |
+
|
24 |
+
def save_reconstruction_images(self, image_latents, new_inv_image_latent, new_G, target_image):
|
25 |
+
if self.save_concatenated_images:
|
26 |
+
save_concat_image(self.concat_base_dir, image_latents, new_inv_image_latent, new_G,
|
27 |
+
self.experiment_creator.old_G,
|
28 |
+
'rec',
|
29 |
+
target_image)
|
30 |
+
|
31 |
+
if self.save_single_images:
|
32 |
+
save_single_image(self.single_base_dir, new_inv_image_latent, new_G, 'rec')
|
33 |
+
target_image.save(f'{self.single_base_dir}/Original.jpg')
|
34 |
+
|
35 |
+
def create_output_dirs(self, full_image_name):
|
36 |
+
output_base_dir_path = f'{paths_config.experiments_output_dir}/{paths_config.input_data_id}/{self.run_id}/{full_image_name}'
|
37 |
+
os.makedirs(output_base_dir_path, exist_ok=True)
|
38 |
+
|
39 |
+
self.concat_base_dir = f'{output_base_dir_path}/concat_images'
|
40 |
+
self.single_base_dir = f'{output_base_dir_path}/single_images'
|
41 |
+
|
42 |
+
os.makedirs(self.concat_base_dir, exist_ok=True)
|
43 |
+
os.makedirs(self.single_base_dir, exist_ok=True)
|
44 |
+
|
45 |
+
def get_image_latent_codes(self, image_name):
|
46 |
+
image_latents = []
|
47 |
+
for method in evaluation_config.evaluated_methods:
|
48 |
+
if method == 'SG2':
|
49 |
+
image_latents.append(torch.load(
|
50 |
+
f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/'
|
51 |
+
f'{paths_config.pti_results_keyword}/{image_name}/0.pt'))
|
52 |
+
else:
|
53 |
+
image_latents.append(torch.load(
|
54 |
+
f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/{method}/{image_name}/0.pt'))
|
55 |
+
new_inv_image_latent = torch.load(
|
56 |
+
f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}/{paths_config.pti_results_keyword}/{image_name}/0.pt')
|
57 |
+
|
58 |
+
return image_latents, new_inv_image_latent
|
59 |
+
|
60 |
+
def save_interfacegan_edits(self, image_latents, new_inv_image_latent, interfacegan_factors, new_G, target_image):
|
61 |
+
new_w_inv_edits = self.latent_editor.get_single_interface_gan_edits(new_inv_image_latent,
|
62 |
+
interfacegan_factors)
|
63 |
+
|
64 |
+
inv_edits = []
|
65 |
+
for latent in image_latents:
|
66 |
+
inv_edits.append(self.latent_editor.get_single_interface_gan_edits(latent, interfacegan_factors))
|
67 |
+
|
68 |
+
for direction, edits in new_w_inv_edits.items():
|
69 |
+
for factor, edit_tensor in edits.items():
|
70 |
+
if self.save_concatenated_images:
|
71 |
+
save_concat_image(self.concat_base_dir, [edits[direction][factor] for edits in inv_edits],
|
72 |
+
new_w_inv_edits[direction][factor],
|
73 |
+
new_G,
|
74 |
+
self.experiment_creator.old_G,
|
75 |
+
f'{direction}_{factor}', target_image)
|
76 |
+
if self.save_single_images:
|
77 |
+
save_single_image(self.single_base_dir, new_w_inv_edits[direction][factor], new_G,
|
78 |
+
f'{direction}_{factor}')
|
79 |
+
|
80 |
+
def save_ganspace_edits(self, image_latents, new_inv_image_latent, factors, new_G, target_image):
|
81 |
+
new_w_inv_edits = self.latent_editor.get_single_ganspace_edits(new_inv_image_latent, factors)
|
82 |
+
inv_edits = []
|
83 |
+
for latent in image_latents:
|
84 |
+
inv_edits.append(self.latent_editor.get_single_ganspace_edits(latent, factors))
|
85 |
+
|
86 |
+
for idx in range(len(new_w_inv_edits)):
|
87 |
+
if self.save_concatenated_images:
|
88 |
+
save_concat_image(self.concat_base_dir, [edit[idx] for edit in inv_edits], new_w_inv_edits[idx],
|
89 |
+
new_G,
|
90 |
+
self.experiment_creator.old_G,
|
91 |
+
f'ganspace_{idx}', target_image)
|
92 |
+
if self.save_single_images:
|
93 |
+
save_single_image(self.single_base_dir, new_w_inv_edits[idx], new_G,
|
94 |
+
f'ganspace_{idx}')
|
95 |
+
|
96 |
+
def run_experiment(self, run_pt, create_other_latents, use_multi_id_training, use_wandb=False):
|
97 |
+
images_counter = 0
|
98 |
+
new_G = None
|
99 |
+
interfacegan_factors = [val / 2 for val in range(-6, 7) if val != 0]
|
100 |
+
ganspace_factors = range(-20, 25, 5)
|
101 |
+
self.experiment_creator.run_experiment(run_pt, create_other_latents, use_multi_id_training, use_wandb)
|
102 |
+
|
103 |
+
if use_multi_id_training:
|
104 |
+
new_G = load_tuned_G(self.run_id, paths_config.multi_id_model_type)
|
105 |
+
|
106 |
+
for idx, image_path in tqdm(enumerate(self.experiment_creator.images_paths),
|
107 |
+
total=len(self.experiment_creator.images_paths)):
|
108 |
+
|
109 |
+
if images_counter >= hyperparameters.max_images_to_invert:
|
110 |
+
break
|
111 |
+
|
112 |
+
image_name = image_path.split('.')[0].split('/')[-1]
|
113 |
+
target_image = Image.open(self.experiment_creator.target_paths[idx])
|
114 |
+
|
115 |
+
if not use_multi_id_training:
|
116 |
+
new_G = load_tuned_G(self.run_id, image_name)
|
117 |
+
|
118 |
+
image_latents, new_inv_image_latent = self.get_image_latent_codes(image_name)
|
119 |
+
|
120 |
+
self.create_output_dirs(image_name)
|
121 |
+
|
122 |
+
self.save_reconstruction_images(image_latents, new_inv_image_latent, new_G, target_image)
|
123 |
+
|
124 |
+
self.save_interfacegan_edits(image_latents, new_inv_image_latent, interfacegan_factors, new_G, target_image)
|
125 |
+
|
126 |
+
self.save_ganspace_edits(image_latents, new_inv_image_latent, ganspace_factors, new_G, target_image)
|
127 |
+
|
128 |
+
target_image.close()
|
129 |
+
torch.cuda.empty_cache()
|
130 |
+
images_counter += 1
|
131 |
+
|
132 |
+
|
133 |
+
def run_pti_and_full_edit(iid):
|
134 |
+
evaluation_config.evaluated_methods = ['SG2Plus', 'e4e', 'SG2']
|
135 |
+
edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
|
136 |
+
run_id=f'{paths_config.input_data_id}_pti_full_edit_{iid}')
|
137 |
+
edit_figure_creator.run_experiment(True, True, use_multi_id_training=False, use_wandb=False)
|
138 |
+
|
139 |
+
|
140 |
+
def pti_no_comparison(iid):
|
141 |
+
evaluation_config.evaluated_methods = []
|
142 |
+
edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
|
143 |
+
run_id=f'{paths_config.input_data_id}_pti_no_comparison_{iid}')
|
144 |
+
edit_figure_creator.run_experiment(True, False, use_multi_id_training=False, use_wandb=False)
|
145 |
+
|
146 |
+
|
147 |
+
def edits_for_existed_experiment(run_id):
|
148 |
+
evaluation_config.evaluated_methods = ['SG2Plus', 'e4e', 'SG2']
|
149 |
+
edit_figure_creator = EditComparison(save_single_images=True, save_concatenated_images=True,
|
150 |
+
run_id=run_id)
|
151 |
+
edit_figure_creator.run_experiment(False, True, use_multi_id_training=False, use_wandb=False)
|
152 |
+
|
153 |
+
|
154 |
+
if __name__ == '__main__':
|
155 |
+
iid = ''.join(choice(ascii_uppercase) for i in range(7))
|
156 |
+
pti_no_comparison(iid)
|
makedirs.py
ADDED
@@ -0,0 +1,83 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import click
|
2 |
+
import os
|
3 |
+
import sys
|
4 |
+
import pickle
|
5 |
+
import numpy as np
|
6 |
+
from PIL import Image
|
7 |
+
import torch
|
8 |
+
from configs import paths_config, hyperparameters, global_config
|
9 |
+
from IPython.display import display
|
10 |
+
import matplotlib.pyplot as plt
|
11 |
+
from scripts.latent_editor_wrapper import LatentEditorWrapper
|
12 |
+
|
13 |
+
|
14 |
+
image_dir_name = '/home/sayantan/processed_images'
|
15 |
+
use_multi_id_training = False
|
16 |
+
global_config.device = 'cuda'
|
17 |
+
paths_config.e4e = '/home/sayantan/PTI/pretrained_models/e4e_ffhq_encode.pt'
|
18 |
+
paths_config.input_data_id = image_dir_name
|
19 |
+
paths_config.input_data_path = f'{image_dir_name}'
|
20 |
+
paths_config.stylegan2_ada_ffhq = '/home/sayantan/PTI/pretrained_models/ffhq.pkl'
|
21 |
+
paths_config.checkpoints_dir = '/home/sayantan/PTI/'
|
22 |
+
paths_config.style_clip_pretrained_mappers = '/home/sayantan/PTI/pretrained_models'
|
23 |
+
hyperparameters.use_locality_regularization = False
|
24 |
+
hyperparameters.lpips_type = 'squeeze'
|
25 |
+
|
26 |
+
model_id = "MYJJDFVGATAT"
|
27 |
+
|
28 |
+
|
29 |
+
|
30 |
+
def display_alongside_source_image(images):
|
31 |
+
res = np.concatenate([np.array(image) for image in images], axis=1)
|
32 |
+
return Image.fromarray(res)
|
33 |
+
|
34 |
+
def load_generators(model_id, image_name):
|
35 |
+
with open(paths_config.stylegan2_ada_ffhq, 'rb') as f:
|
36 |
+
old_G = pickle.load(f)['G_ema'].cuda()
|
37 |
+
|
38 |
+
with open(f'{paths_config.checkpoints_dir}/model_{model_id}_{image_name}.pt', 'rb') as f_new:
|
39 |
+
new_G = torch.load(f_new).cuda()
|
40 |
+
|
41 |
+
return old_G, new_G
|
42 |
+
|
43 |
+
def plot_syn_images(syn_images,text):
|
44 |
+
for img in syn_images:
|
45 |
+
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
|
46 |
+
plt.axis('off')
|
47 |
+
resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
|
48 |
+
display(resized_image)
|
49 |
+
#wandb.log({text: [wandb.Image(resized_image, caption="Label")]})
|
50 |
+
del img
|
51 |
+
del resized_image
|
52 |
+
torch.cuda.empty_cache()
|
53 |
+
|
54 |
+
def syn_images_wandb(img):
|
55 |
+
img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).detach().cpu().numpy()[0]
|
56 |
+
plt.axis('off')
|
57 |
+
resized_image = Image.fromarray(img,mode='RGB').resize((256,256))
|
58 |
+
return resized_image
|
59 |
+
|
60 |
+
@click.command()
|
61 |
+
@click.pass_context
|
62 |
+
@click.option('--image_name', prompt='image name', help='The name for image')
|
63 |
+
|
64 |
+
def makedir(ctx: click.Context,image_name):
|
65 |
+
generator_type = paths_config.multi_id_model_type if use_multi_id_training else image_name
|
66 |
+
old_G, new_G = load_generators(model_id, generator_type)
|
67 |
+
w_path_dir = f'{paths_config.embedding_base_dir}/{paths_config.input_data_id}'
|
68 |
+
|
69 |
+
embedding_dir = f'{w_path_dir}/{paths_config.pti_results_keyword}/{image_name}'
|
70 |
+
w_pivot = torch.load(f'{embedding_dir}/0.pt')
|
71 |
+
|
72 |
+
old_image = old_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
|
73 |
+
new_image = new_G.synthesis(w_pivot, noise_mode='const', force_fp32 = True)
|
74 |
+
|
75 |
+
latent_editor = LatentEditorWrapper()
|
76 |
+
latents_after_edit = latent_editor.get_single_interface_gan_edits(w_pivot, [i for i in range(-5,5)])
|
77 |
+
|
78 |
+
for direction, factor_and_edit in latents_after_edit.items():
|
79 |
+
for editkey in factor_and_edit.keys():
|
80 |
+
os.makedirs(f"/home/sayantan/PTI/{direction}/{editkey}")
|
81 |
+
|
82 |
+
if __name__ == '__main__':
|
83 |
+
makedir()
|
models/StyleCLIP/__init__.py
ADDED
File without changes
|
models/StyleCLIP/criteria/__init__.py
ADDED
File without changes
|
models/StyleCLIP/criteria/clip_loss.py
ADDED
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
|
2 |
+
import torch
|
3 |
+
import clip
|
4 |
+
|
5 |
+
|
6 |
+
class CLIPLoss(torch.nn.Module):
|
7 |
+
|
8 |
+
def __init__(self, opts):
|
9 |
+
super(CLIPLoss, self).__init__()
|
10 |
+
self.model, self.preprocess = clip.load("ViT-B/32", device="cuda")
|
11 |
+
self.upsample = torch.nn.Upsample(scale_factor=7)
|
12 |
+
self.avg_pool = torch.nn.AvgPool2d(kernel_size=opts.stylegan_size // 32)
|
13 |
+
|
14 |
+
def forward(self, image, text):
|
15 |
+
image = self.avg_pool(self.upsample(image))
|
16 |
+
similarity = 1 - self.model(image, text)[0] / 100
|
17 |
+
return similarity
|
models/StyleCLIP/criteria/id_loss.py
ADDED
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import torch
|
2 |
+
from torch import nn
|
3 |
+
|
4 |
+
from models.facial_recognition.model_irse import Backbone
|
5 |
+
|
6 |
+
|
7 |
+
class IDLoss(nn.Module):
|
8 |
+
def __init__(self, opts):
|
9 |
+
super(IDLoss, self).__init__()
|
10 |
+
print('Loading ResNet ArcFace')
|
11 |
+
self.facenet = Backbone(input_size=112, num_layers=50, drop_ratio=0.6, mode='ir_se')
|
12 |
+
self.facenet.load_state_dict(torch.load(opts.ir_se50_weights))
|
13 |
+
self.pool = torch.nn.AdaptiveAvgPool2d((256, 256))
|
14 |
+
self.face_pool = torch.nn.AdaptiveAvgPool2d((112, 112))
|
15 |
+
self.facenet.eval()
|
16 |
+
self.opts = opts
|
17 |
+
|
18 |
+
def extract_feats(self, x):
|
19 |
+
if x.shape[2] != 256:
|
20 |
+
x = self.pool(x)
|
21 |
+
x = x[:, :, 35:223, 32:220] # Crop interesting region
|
22 |
+
x = self.face_pool(x)
|
23 |
+
x_feats = self.facenet(x)
|
24 |
+
return x_feats
|
25 |
+
|
26 |
+
def forward(self, y_hat, y):
|
27 |
+
n_samples = y.shape[0]
|
28 |
+
y_feats = self.extract_feats(y) # Otherwise use the feature from there
|
29 |
+
y_hat_feats = self.extract_feats(y_hat)
|
30 |
+
y_feats = y_feats.detach()
|
31 |
+
loss = 0
|
32 |
+
sim_improvement = 0
|
33 |
+
count = 0
|
34 |
+
for i in range(n_samples):
|
35 |
+
diff_target = y_hat_feats[i].dot(y_feats[i])
|
36 |
+
loss += 1 - diff_target
|
37 |
+
count += 1
|
38 |
+
|
39 |
+
return loss / count, sim_improvement / count
|
models/StyleCLIP/global_directions/GUI.py
ADDED
@@ -0,0 +1,103 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
|
2 |
+
|
3 |
+
from tkinter import Tk,Frame ,Label,Button,messagebox,Canvas,Text,Scale
|
4 |
+
from tkinter import HORIZONTAL
|
5 |
+
|
6 |
+
class View():
|
7 |
+
def __init__(self,master):
|
8 |
+
|
9 |
+
self.width=600
|
10 |
+
self.height=600
|
11 |
+
|
12 |
+
|
13 |
+
self.root=master
|
14 |
+
self.root.geometry("600x600")
|
15 |
+
|
16 |
+
self.left_frame=Frame(self.root,width=600)
|
17 |
+
self.left_frame.pack_propagate(0)
|
18 |
+
self.left_frame.pack(fill='both', side='left', expand='True')
|
19 |
+
|
20 |
+
self.retrieval_frame=Frame(self.root,bg='snow3')
|
21 |
+
self.retrieval_frame.pack_propagate(0)
|
22 |
+
self.retrieval_frame.pack(fill='both', side='right', expand='True')
|
23 |
+
|
24 |
+
self.bg_frame=Frame(self.left_frame,bg='snow3',height=600,width=600)
|
25 |
+
self.bg_frame.pack_propagate(0)
|
26 |
+
self.bg_frame.pack(fill='both', side='top', expand='True')
|
27 |
+
|
28 |
+
self.command_frame=Frame(self.left_frame,bg='snow3')
|
29 |
+
self.command_frame.pack_propagate(0)
|
30 |
+
self.command_frame.pack(fill='both', side='bottom', expand='True')
|
31 |
+
# self.command_frame.grid(row=1, column=0,padx=0, pady=0)
|
32 |
+
|
33 |
+
self.bg=Canvas(self.bg_frame,width=self.width,height=self.height, bg='gray')
|
34 |
+
self.bg.place(relx=0.5, rely=0.5, anchor='center')
|
35 |
+
|
36 |
+
self.mani=Canvas(self.retrieval_frame,width=1024,height=1024, bg='gray')
|
37 |
+
self.mani.grid(row=0, column=0,padx=0, pady=42)
|
38 |
+
|
39 |
+
self.SetCommand()
|
40 |
+
|
41 |
+
|
42 |
+
|
43 |
+
|
44 |
+
def run(self):
|
45 |
+
self.root.mainloop()
|
46 |
+
|
47 |
+
def helloCallBack(self):
|
48 |
+
category=self.set_category.get()
|
49 |
+
messagebox.showinfo( "Hello Python",category)
|
50 |
+
|
51 |
+
def SetCommand(self):
|
52 |
+
|
53 |
+
tmp = Label(self.command_frame, text="neutral", width=10 ,bg='snow3')
|
54 |
+
tmp.grid(row=1, column=0,padx=10, pady=10)
|
55 |
+
|
56 |
+
tmp = Label(self.command_frame, text="a photo of a", width=10 ,bg='snow3')
|
57 |
+
tmp.grid(row=1, column=1,padx=10, pady=10)
|
58 |
+
|
59 |
+
self.neutral = Text ( self.command_frame, height=2, width=30)
|
60 |
+
self.neutral.grid(row=1, column=2,padx=10, pady=10)
|
61 |
+
|
62 |
+
|
63 |
+
tmp = Label(self.command_frame, text="target", width=10 ,bg='snow3')
|
64 |
+
tmp.grid(row=2, column=0,padx=10, pady=10)
|
65 |
+
|
66 |
+
tmp = Label(self.command_frame, text="a photo of a", width=10 ,bg='snow3')
|
67 |
+
tmp.grid(row=2, column=1,padx=10, pady=10)
|
68 |
+
|
69 |
+
self.target = Text ( self.command_frame, height=2, width=30)
|
70 |
+
self.target.grid(row=2, column=2,padx=10, pady=10)
|
71 |
+
|
72 |
+
tmp = Label(self.command_frame, text="strength", width=10 ,bg='snow3')
|
73 |
+
tmp.grid(row=3, column=0,padx=10, pady=10)
|
74 |
+
|
75 |
+
self.alpha = Scale(self.command_frame, from_=-15, to=25, orient=HORIZONTAL,bg='snow3', length=250,resolution=0.01)
|
76 |
+
self.alpha.grid(row=3, column=2,padx=10, pady=10)
|
77 |
+
|
78 |
+
|
79 |
+
tmp = Label(self.command_frame, text="disentangle", width=10 ,bg='snow3')
|
80 |
+
tmp.grid(row=4, column=0,padx=10, pady=10)
|
81 |
+
|
82 |
+
self.beta = Scale(self.command_frame, from_=0.08, to=0.4, orient=HORIZONTAL,bg='snow3', length=250,resolution=0.001)
|
83 |
+
self.beta.grid(row=4, column=2,padx=10, pady=10)
|
84 |
+
|
85 |
+
self.reset = Button(self.command_frame, text='Reset')
|
86 |
+
self.reset.grid(row=5, column=1,padx=10, pady=10)
|
87 |
+
|
88 |
+
|
89 |
+
self.set_init = Button(self.command_frame, text='Accept')
|
90 |
+
self.set_init.grid(row=5, column=2,padx=10, pady=10)
|
91 |
+
|
92 |
+
#%%
|
93 |
+
if __name__ == "__main__":
|
94 |
+
master=Tk()
|
95 |
+
self=View(master)
|
96 |
+
self.run()
|
97 |
+
|
98 |
+
|
99 |
+
|
100 |
+
|
101 |
+
|
102 |
+
|
103 |
+
|
models/StyleCLIP/global_directions/GenerateImg.py
ADDED
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
|
2 |
+
import os
|
3 |
+
import numpy as np
|
4 |
+
import argparse
|
5 |
+
from manipulate import Manipulator
|
6 |
+
|
7 |
+
from PIL import Image
|
8 |
+
#%%
|
9 |
+
|
10 |
+
if __name__ == "__main__":
|
11 |
+
parser = argparse.ArgumentParser(description='Process some integers.')
|
12 |
+
|
13 |
+
parser.add_argument('--dataset_name',type=str,default='ffhq',
|
14 |
+
help='name of dataset, for example, ffhq')
|
15 |
+
|
16 |
+
args = parser.parse_args()
|
17 |
+
dataset_name=args.dataset_name
|
18 |
+
|
19 |
+
if not os.path.isdir('./data/'+dataset_name):
|
20 |
+
os.system('mkdir ./data/'+dataset_name)
|
21 |
+
#%%
|
22 |
+
M=Manipulator(dataset_name=dataset_name)
|
23 |
+
np.set_printoptions(suppress=True)
|
24 |
+
print(M.dataset_name)
|
25 |
+
#%%
|
26 |
+
|
27 |
+
M.img_index=0
|
28 |
+
M.num_images=50
|
29 |
+
M.alpha=[0]
|
30 |
+
M.step=1
|
31 |
+
lindex,bname=0,0
|
32 |
+
|
33 |
+
M.manipulate_layers=[lindex]
|
34 |
+
codes,out=M.EditOneC(bname)
|
35 |
+
#%%
|
36 |
+
|
37 |
+
for i in range(len(out)):
|
38 |
+
img=out[i,0]
|
39 |
+
img=Image.fromarray(img)
|
40 |
+
img.save('./data/'+dataset_name+'/'+str(i)+'.jpg')
|
41 |
+
#%%
|
42 |
+
w=np.load('./npy/'+dataset_name+'/W.npy')
|
43 |
+
|
44 |
+
tmp=w[:M.num_images]
|
45 |
+
tmp=tmp[:,None,:]
|
46 |
+
tmp=np.tile(tmp,(1,M.Gs.components.synthesis.input_shape[1],1))
|
47 |
+
|
48 |
+
np.save('./data/'+dataset_name+'/w_plus.npy',tmp)
|
49 |
+
|
50 |
+
|