Upload app.py

app.py

@@ -1,139 +1,77 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Tue Apr 26 21:02:31 2022
+
+@author: pc
+"""
 
-import os
 import pickle
-import sys
-import subprocess
-import imageio
 import numpy as np
-import scipy.interpolate
 import torch
-from tqdm import tqdm
 import gradio as gr 
-
+import sys
+import subprocess
+import os
+from typing import Tuple
+import PIL.Image
 
 os.system("git clone https://github.com/NVlabs/stylegan3")
 
 sys.path.append("stylegan3")
 
-def layout_grid(img, grid_w=None, grid_h=1, float_to_uint8=True, chw_to_hwc=True, to_numpy=True):
-    batch_size, channels, img_h, img_w = img.shape
-    if grid_w is None:
-        grid_w = batch_size // grid_h
-    assert batch_size == grid_w * grid_h
-    if float_to_uint8:
-        img = (img * 127.5 + 128).clamp(0, 255).to(torch.uint8)
-    img = img.reshape(grid_h, grid_w, channels, img_h, img_w)
-    img = img.permute(2, 0, 3, 1, 4)
-    img = img.reshape(channels, grid_h * img_h, grid_w * img_w)
-    if chw_to_hwc:
-        img = img.permute(1, 2, 0)
-    if to_numpy:
-        img = img.cpu().numpy()
-    return img
 
 
+def make_transform(translate: Tuple[float,float], angle: float):
+    m = np.eye(3)
+    s = np.sin(angle/360.0*np.pi*2)
+    c = np.cos(angle/360.0*np.pi*2)
+    m[0][0] = c
+    m[0][1] = s
+    m[0][2] = translate[0]
+    m[1][0] = -s
+    m[1][1] = c
+    m[1][2] = translate[1]
+    return m
+
 
 
 network_pkl='braingan-400.pkl'
 with open(network_pkl, 'rb') as f:
     G = pickle.load(f)['G_ema'] 
-    
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 G.eval()
 G.to(device)
-def predict(Seed,choices,choices2):
-
-  shuffle_seed=None
-  kind='cubic' 
-  num_keyframes=None
-  wraps=2
-  psi=0.5 
-  s1=Seed
-
-  
-  if choices=='4x2':
-    grid_w = 4
-    grid_h = 2
-    if choices2=="Large Video":
-        seeds=(np.arange(s1-16,s1)).tolist()  
-        w_frames=60*4
-    if choices2=="Small Video":
-        seeds=(np.arange(s1-8,s1)).tolist()  
-        w_frames=30*4
-        
-  if choices=='2x1':
-    grid_w = 2
-    grid_h = 1
-    if choices2=="Large Video":
-        seeds=(np.arange(s1-4,s1)).tolist()  
-        w_frames=60*4
-    if choices2=="Small Video":
-        seeds=(np.arange(s1-2,s1)).tolist()  
-        w_frames=30*4
-
-
-
-  mp4='ex.mp4'
-  truncation_psi=1
-  num_keyframes=None
-
-
-  if num_keyframes is None:
-      if len(seeds) % (grid_w*grid_h) != 0:
-          raise ValueError('Number of input seeds must be divisible by grid W*H')
-      num_keyframes = len(seeds) // (grid_w*grid_h)
-
-  all_seeds = np.zeros(num_keyframes*grid_h*grid_w, dtype=np.int64)
-  for idx in range(num_keyframes*grid_h*grid_w):
-      all_seeds[idx] = seeds[idx % len(seeds)]
-
-  if shuffle_seed is not None:
-      rng = np.random.RandomState(seed=shuffle_seed)
-      rng.shuffle(all_seeds)
-
-  zs = torch.from_numpy(np.stack([np.random.RandomState(seed).randn(G.z_dim) for seed in all_seeds])).to(device)
-  ws = G.mapping(z=zs, c=None, truncation_psi=psi)
-  _ = G.synthesis(ws[:1]) # warm up
-  ws = ws.reshape(grid_h, grid_w, num_keyframes, *ws.shape[1:])
-
-  # Interpolation.
-  grid = []
-  for yi in range(grid_h):
-      row = []
-      for xi in range(grid_w):
-          x = np.arange(-num_keyframes * wraps, num_keyframes * (wraps + 1))
-          y = np.tile(ws[yi][xi].cpu().numpy(), [wraps * 2 + 1, 1, 1])
-          interp = scipy.interpolate.interp1d(x, y, kind=kind, axis=0)
-          row.append(interp)
-      grid.append(row)
-
-  # Render video.
-  video_out = imageio.get_writer(mp4, mode='I', fps=30, codec='libx264')
-  for frame_idx in tqdm(range(num_keyframes * w_frames)):
-      imgs = []
-      for yi in range(grid_h):
-          for xi in range(grid_w):
-              interp = grid[yi][xi]
-              w = torch.from_numpy(interp(frame_idx / w_frames)).to(device)
-              img = G.synthesis(ws=w.unsqueeze(0), noise_mode='const')[0]
-              imgs.append(img)
-      video_out.append_data(layout_grid(torch.stack(imgs), grid_w=grid_w, grid_h=grid_h))
-  video_out.close()
-  return mp4
-
-
-
-choices=['4x2','2x1']
-choices2=["Large Video","Small Video"]
 
+def predict(Seed,noise_mode,truncation_psi,trans_x,trans_y,angle):
+
+  # Generate images.
+    z = torch.from_numpy(np.random.RandomState(Seed).randn(1, G.z_dim)).to(device)
+    label = torch.zeros([1, G.c_dim], device=device)
+    # Construct an inverse rotation/translation matrix and pass to the generator.  The
+    # generator expects this matrix as an inverse to avoid potentially failing numerical
+    # operations in the network.
+    if hasattr(G.synthesis, 'input'):
+        m = make_transform((trans_x,trans_y), angle)
+        m = np.linalg.inv(m)
+        G.synthesis.input.transform.copy_(torch.from_numpy(m))
+
+    img = G(z, label, truncation_psi=truncation_psi, noise_mode=noise_mode)
+    img = (img.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8)
+      
+    return PIL.Image.fromarray(img[0].cpu().numpy()[:,:,0])
+
+
+
+noises=['const', 'random', 'none']
 interface=gr.Interface(fn=predict, title="Brain MR Image Generation with StyleGAN-2",
                        description = "",
                        article = "Author: S.Serdar Helli",
-                       inputs=[gr.inputs.Slider( minimum=16, maximum=2**10,label='Seed'),gr.inputs.Radio( choices=choices,  default='4x2',label='Image Grid'),
-                       gr.inputs.Radio( choices=choices2,  default="Small Video",label='Video Size - It depends on usage of cuda')],
-                       outputs=gr.outputs.Video(label='Video'),
-                       live=False)
-                                              
+                       inputs=[gr.inputs.Slider( minimum=0, maximum=2**10,label='Seed'),gr.inputs.Radio( choices=noises,  default='const',label='Noise Mods'),
+                                           gr.inputs.Slider(0, 2, step=0.05, default=1, label='Truncation psi'),
+                                           gr.inputs.Slider(-1, 1, step=0.05, default=0, label='Translate X'),
+                                           gr.inputs.Slider(-1, 1, step=0.05, default=0, label='Translate Y'),
+                                           gr.inputs.Slider(-180, 180, step=5, default=0, label='Angle'),],
+                       outputs=gr.outputs.Image( type="numpy", label="Output"))
 
 
-interface.launch(debug=True,show_error=True)
+interface.launch(debug=True)