File size: 1,718 Bytes
b1c8049 |
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 |
import os
from PIL import Image
from torchvision import transforms as T
from torchvision.transforms import Compose, Resize, ToTensor, Normalize, RandomCrop, RandomHorizontalFlip
from torchvision.utils import make_grid
from torch.utils.data import DataLoader
from huggan.pytorch.cyclegan.modeling_cyclegan import GeneratorResNet
import torch.nn as nn
import torch
import gradio as gr
from collections import OrderedDict
import glob
def pred_pipeline(img, transforms):
orig_shape = img.shape
input = transforms(img)
input = input.unsqueeze(0)
output = model(input)
out_img = make_grid(output,#.detach().cpu(),
nrow=1, normalize=True)
out_transform = Compose([
T.Resize(orig_shape[:2]),
T.ToPILImage()
])
return out_transform(out_img)
n_channels = 3
image_size = 512
input_shape = (image_size, image_size)
transform = Compose([
T.ToPILImage(),
T.Resize(input_shape),
ToTensor(),
Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])
model = GeneratorResNet.from_pretrained('Chris1/sim2real', input_shape=(n_channels, image_size, image_size),
num_residual_blocks=9)
gr.Interface(lambda image: pred_pipeline(image, transform),
inputs=gr.inputs.Image( label='input synthetic image'),
outputs=gr.outputs.Image( type="pil",label='style transfer to the real world'),#plot,
title = "GTA5(simulated) to Cityscapes (real) translation",
examples = [
[example] for example in glob.glob('./samples/*.png')
])\
.launch()
#iface = gr.Interface(fn=greet, inputs="text", outputs="text")
#iface.launch()
|