muratcelik/Image_Inpainting_w_context-encoder

1import streamlit as st2 3from model import get_mask4from model import ContextGenerator5 6import torch7from torchvision import transforms8from PIL import Image9import random10import torchvision.transforms as T11 12 13 14img_size = 12815out_size = 12816mask_size = 3217 18transform = transforms.Compose(19        [20            transforms.Resize((img_size, img_size)),21            transforms.ToTensor()22                                23        ]24    )25 26 27generator = ContextGenerator(bottleneck_dim=4000, img_size=img_size, out_size=out_size)28generator.load_state_dict(torch.load("generato.pt",map_location=torch.device("cpu")))29generator.eval()30 31st.header("Image Inpainting")32st.write("This project is based on the article Context Encoders: Feature Learning by Inpainting. \n\n" + 33"https://github.com/fbuchert/context-encoder-pytorch This repo has been used.")34 35 36img_file_buffer = st.file_uploader('Upload a PNG image', type=['png', 'jpg'])37if img_file_buffer is not None:38    img = Image.open(img_file_buffer)39    x,y = random.randint(25,128-64),random.randint(25,128-64)40    img_transform = transform(img)41    masked_samples, true_masked_part, mask_coordinates = get_mask(img_transform, mask_size,[x,y])42    outG = generator(masked_samples.unsqueeze(0))43    masked_samples_result = masked_samples.clone()44    masked_samples_result[:, x: x + mask_size, y: y + mask_size] = outG[0][:, x: x + mask_size, y: y + mask_size]45 46    col1, col2, col3 = st.columns(3)47    with col1:48        st.header("Original")49        st.image(img_transform.permute(1, 2, 0).cpu().detach().numpy())50 51    with col2:52        st.header("Masked")53        st.image(masked_samples.permute(1, 2, 0).cpu().detach().numpy(), clamp=True, channels='RGB')54 55    with col3:56        st.header("Result")57        st.image(masked_samples_result.permute(1, 2, 0).cpu().detach().numpy(), clamp=True, channels='RGB')58