added generated file

Dockerfile

@@ -5,7 +5,13 @@ WORKDIR /data
 COPY requirements.txt ./
 RUN pip install --no-cache-dir -r requirements.txt
 
-EXPOSE 7860
+RUN apt-get update && \
+    apt-get install -y libglib2.0-0 && \
+    apt-get install libgl1-mesa-glx && \
+    apt-get install -y ffmpeg && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
+EXPOSE 7860
 COPY . .
 CMD ["streamlit", "run", "app.py", "--server.port", "7860"]

app.py

@@ -1,12 +1,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 cv2
+import os
+import subprocess
 
 def check_gpu():
     return tf.test.is_gpu_available(cuda_only=False, min_cuda_compute_capability=None)
 
-st.title('GPU Availability Checker')
+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
 
-if check_gpu():
-    st.success('GPU is available!')
-else:
-    st.warning('GPU is not available.')
+# 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)
+
+fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for .mp4
+temp_video_path="output.mp4"
+out = cv2.VideoWriter(temp_video_path, fourcc, 20.0, (256, 256), isColor=False)
+
+for i in range(vid.shape[0]):
+    frame = vid[i]
+    out.write(frame)
+out.release()
+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)