|
|
import pickle |
|
|
import imageio |
|
|
import numpy as np |
|
|
import scipy.interpolate |
|
|
import torch |
|
|
from tqdm import tqdm |
|
|
import gradio as gr |
|
|
from huggingface_hub import hf_hub_download |
|
|
import sys |
|
|
import os |
|
|
import subprocess |
|
|
|
|
|
|
|
|
subprocess.Popen('git clone https://huggingface.co/SerdarHelli/BrainMRIGAN') |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
sys.path.insert(0, 'Brain-MR-Image-Generation-GAN ') |
|
|
|
|
|
|
|
|
network_pkl='braingan-400.pkl' |
|
|
with open(network_pkl, 'rb') as f: |
|
|
G = pickle.load(f)['G_ema'] |
|
|
|
|
|
def predict(Seed,choices): |
|
|
device = torch.device('cuda') |
|
|
G.eval() |
|
|
G.to(device) |
|
|
shuffle_seed=None |
|
|
w_frames=60*4 |
|
|
kind='cubic' |
|
|
num_keyframes=None |
|
|
wraps=2 |
|
|
psi=1 |
|
|
device=torch.device('cuda') |
|
|
|
|
|
|
|
|
if choices=='4x2': |
|
|
grid_w = 4 |
|
|
grid_h = 2 |
|
|
s1=Seed |
|
|
seeds=(np.arange(s1-16,s1)).tolist() |
|
|
if choices=='2x1': |
|
|
grid_w = 2 |
|
|
grid_h = 1 |
|
|
s1=Seed |
|
|
seeds=(np.arange(s1-4,s1)).tolist() |
|
|
|
|
|
|
|
|
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]) |
|
|
ws = ws.reshape(grid_h, grid_w, num_keyframes, *ws.shape[1:]) |
|
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
|
video_out = imageio.get_writer(mp4, mode='I', fps=60, 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 'ex.mp4' |
|
|
|
|
|
|
|
|
|
|
|
choices=['4x2','2x1'] |
|
|
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')], |
|
|
outputs=gr.outputs.Video(label='Video')) |
|
|
|
|
|
|
|
|
interface.launch(debug=True) |
