File size: 2,083 Bytes
d813284 53b5c44 d813284 7a6408f d813284 7a6408f d813284 |
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 |
import gradio as gr
from huggingface_hub import hf_hub_download
import torch
from networks import define_G
from torchvision import transforms
REPO_ID = "Launchpad/ditto"
FILENAME = "model.pth"
model_dict = torch.load(
hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
)
generator = define_G(input_nc=3, output_nc=3, ngf=64, netG="resnet_9blocks", norm="instance")
generator.load_state_dict(model_dict)
generator.eval()
# set up transforms for model
encode = transforms.Compose([
transforms.ToTensor(),
transforms.Resize((256, 256))
])
transform = transforms.ToPILImage()
def generate_pokemon(pet_img):
# encode image
encoded_img = encode(pet_img)
# evaluate model on pet image
with torch.no_grad():
generated_img = generator(encoded_img)
# transform to PIL image
return transform(generated_img)
with gr.Blocks() as demo:
with gr.Row():
with gr.Column(scale=1):
gr.Image("https://www.ocf.berkeley.edu/~launchpad/media/uploads/project_logos/Ditto.png", elem_id="logo-img", show_label=False, show_share_button=False, show_download_button=False)
with gr.Column(scale=3):
gr.Markdown("""Ditto is a [Launchpad](https://launchpad.studentorg.berkeley.edu/) project (Fall 2022) that transfers styles of Pokemon sprites onto pet images using GANs and contrastive learning.
<br/><br/>
**Model**: [ditto](https://huggingface.co/Launchpad/ditto)
<br/>
**Developed by**: Kiran Suresh, Annie Lee, Chloe Wong, Tony Xin, Sebastian Zhao
<br/>
**Examples**: [Oxford-IIIT Pet Dataset](https://www.robots.ox.ac.uk/~vgg/data/pets/)
"""
)
gr.Interface(fn=generate_pokemon,
inputs=gr.Image(),
outputs="image",
examples=["data/german_shorthaired_164.jpg", "data/samoyed_189.jpg", "data/shiba_inu_139.jpg"]
)
if __name__ == '__main__':
demo.launch()
|