Spaces:
Sleeping
Sleeping
File size: 2,184 Bytes
ae0af75 |
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 |
import torch
import numpy as np
import gradio as gr
import os
from PIL import Image
import torchvision.transforms as T
from config.core import config
from utility.helper import load_model_weights, init_generator_model
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = init_generator_model()
model = load_model_weights(
model=model,
checkpoint_path=config.CKPT_PATH,
device=device,
prefix="gen",
)
# Transformation
transform_face = T.Compose([
T.Resize((config.IMAGE_SIZE, config.IMAGE_SIZE)),
T.ToTensor(),
T.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))
])
def inference(image: Image):
# transforms the target image and add a batch dimension
img = transform_face(image)
img_un = img.unsqueeze(0)
image_transform = img_un * 0.5 + 0.5 # Normalize from Tanh
im_detach = image_transform.detach().cpu().squeeze(0)
im_permute = im_detach.permute(1, 2, 0)
im_array = im_permute.numpy()
# Scale values to 0-255 range
im_array = (im_array * 255).astype(np.uint8)
# Convert numpy array to PIL Image
im_pil = Image.fromarray(im_array)
# Inference the image
model.eval()
with torch.inference_mode():
c2f = model(img_un)
c2f = c2f * 0.5 + 0.5 # Normalize from Tanh
image_unflat = c2f.detach().cpu().squeeze(0) # Remove batch dimension
image = image_unflat.permute(1, 2, 0) # Permute to (H, W, C)
# Convert image to numpy array
image_array = image.numpy()
# Scale values to 0-255 range
image_array = (image_array * 255).astype(np.uint8)
# Convert numpy array to PIL Image
image = Image.fromarray(image_array)
return im_pil, image
demo = gr.Interface(
fn=inference,
inputs=gr.Image(type="pil"),
outputs=[
gr.Image(label="Original after Transform"),
gr.Image(label="Converted by Model")
],
title="Pix2Pix Face to Comic",
description="A implementation Pix2Pix from Scratch Pytorch",
examples=[f"data/examples/{i}" for i in os.listdir("data/examples") if i.endswith(('.png', '.jpg', '.jpeg', '.gif'))]
)
demo.launch()
|