Spaces:
Sleeping
Sleeping
File size: 2,817 Bytes
aa4251c a47a354 92a78e5 a47a354 aa4251c a47a354 aa4251c a47a354 92a78e5 a47a354 92a78e5 a47a354 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 |
import streamlit as st
import tensorflow as tf
import pickle
import numpy as np
from pathlib import Path
import dnnlib
from dnnlib import tflib
import imageio
import os
import subprocess
def check_gpu():
return tf.test.is_gpu_available(cuda_only=False, min_cuda_compute_capability=None)
model_path = 'best_net.pkl'
#define load model functions
_cached_networks = dict()
def load_networks(path):
if path in _cached_networks:
return _cached_networks[path]
stream = open(path, 'rb')
tflib.init_tf()
with stream:
G, D, Gs = pickle.load(stream, encoding='latin1')
_cached_networks[path] = G, D, Gs
return G, D, Gs
# Code to load the StyleGAN2 Model
def load_model():
_G, _D, Gs = load_networks(model_path)
noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
Gs_kwargs = dnnlib.EasyDict()
Gs_kwargs.output_transform = dict(func=tflib.convert_images_to_uint8, nchw_to_nhwc=True)
Gs_kwargs.randomize_noise = False
return Gs, noise_vars, Gs_kwargs
#define helper functions
def get_control_latent_vectors(path):
files = [x for x in Path(path).iterdir() if str(x).endswith('.npy')]
latent_vectors = {f.name[:-4]:np.load(f) for f in files}
return latent_vectors
#load latent directions
latent_controls = get_control_latent_vectors('trajectories/')
def generate_image_from_projected_latents(latent_vector):
images = Gs.components.synthesis.run(latent_vector, **Gs_kwargs)
return images
def frame_to_frame(latent_code):
modified_latent_code = np.copy(latent_code)
full_video = [generate_image_from_projected_latents(modified_latent_code)]
for i in range(49):
modified_latent_code = modified_latent_code + latent_controls[f'{i}{i+1}']
ims = generate_image_from_projected_latents(modified_latent_code)
full_video.append(ims)
return np.array(full_video).squeeze()
#load the model
Gs, noise_vars, Gs_kwargs = load_model()
#select a random latent code
rnd = np.random.RandomState(3)
z = rnd.randn(1, *Gs.input_shape[1:])
noise_vars = [var for name, var in Gs.components.synthesis.vars.items() if name.startswith('noise')]
tflib.set_vars({var: rnd.randn(*var.shape.as_list()) for var in noise_vars})
random_img_latent_code = Gs.components.mapping.run(z,None)
#make it be ED frame
random_img_latent_code -= 0.7*latent_controls['time']
vid = frame_to_frame(random_img_latent_code)
temp_video_path="output.mp4"
writer=imageio.get_writer(temp_video_path, fps=20)
for i in range(vid.shape[0]):
frame = vid[i]
writer.append_data(frame)
writer.close()
out_path = "fixed_out.mp4"
command = ["ffmpeg", "-i", temp_video_path, "-vcodec", "libx264", out_path]
subprocess.run(command)
st.video(out_path)
os.remove(temp_video_path)
os.remove(out_path) |