<jupyter_start><jupyter_text>Neural Processes for ImagesThis notebook contains examples of Neural Processes for images and how these can be used for various tasks like inpainting.  Load a trained model<jupyter_code>import json
from neural_process import NeuralProcessImg

# Load config file for celebA model
folder = 'trained_models/celeba'
config_file = folder + '/config.json'
model_file = folder + '/model.pt'

with open(config_file) as f:
    config = json.load(f)

# Load trained model
model = NeuralProcessImg(config["img_size"], 
                         config["r_dim"], 
                         config["h_dim"], 
                         config["z_dim"]).to(device)
model.load_state_dict(torch.load(model_file, map_location=lambda storage, loc: storage))<jupyter_output><empty_output><jupyter_text>Visualize some CelebA samples<jupyter_code>import imageio
from torchvision.utils import make_grid

# Read images into torch.Tensor
all_imgs = torch.zeros(8, 3, 32, 32)

for i in range(8):
    img = imageio.imread('imgs/celeba-examples/{}.png'.format(i + 1))
    all_imgs[i] = torch.Tensor(img.transpose(2, 0, 1) / 255.)

# Visualize sample on a grid
img_grid = make_grid(all_imgs, nrow=4, pad_value=1.)
plt.imshow(img_grid.permute(1, 2, 0).numpy())<jupyter_output><empty_output><jupyter_text>Inpainting images with Neural ProcessesInpainting is the task of inferring missing pixels in a partially occluded image. Here we show examples of how Neural Processes can be used to solve this problem.  Occluding image<jupyter_code># Select one of the images to perform inpainting
img = all_imgs[0]

# Define a binary mask to occlude image. For Neural Processes,
# the context points will be defined as the visible pixels
context_mask = torch.zeros((32, 32)).byte()
context_mask[:16, :] = 1  # Top half of pixels are visible

# Show occluded image
occluded_img = img * context_mask.float()
plt.imshow(occluded_img.permute(1, 2, 0).numpy())<jupyter_output><empty_output><jupyter_text>Generating inpaintings<jupyter_code>from utils import inpaint

num_inpaintings = 8  # Number of inpaintings to sample from model
all_inpaintings = torch.zeros(num_inpaintings, 3, 32, 32)

# Sample several inpaintings
for i in range(num_inpaintings):
    all_inpaintings[i] = inpaint(model, img, context_mask, device)

# Visualize inpainting results on a grid
inpainting_grid = make_grid(all_inpaintings, nrow=4, pad_value=1.)
plt.imshow(inpainting_grid.permute(1, 2, 0).numpy())<jupyter_output><empty_output><jupyter_text>As can be seen, the inpaintings match the context pixels and are fairly diverse.  Different masksWe can use a variety of masks and image to test the model.<jupyter_code># Select one of the images to perform inpainting
img = all_imgs[1]

# Define a random mask
context_mask = torch.Tensor(32, 32).uniform_() > 0.9

# Visualize occluded image
occluded_img = img * context_mask.float()
plt.imshow(occluded_img.permute(1, 2, 0).numpy())
num_inpaintings = 8  # Number of inpaintings to sample from model
all_inpaintings = torch.zeros(num_inpaintings, 3, 32, 32)

# Sample several inpaintings
for i in range(num_inpaintings):
    all_inpaintings[i] = inpaint(model, img, context_mask, device)

# Visualize inpainting results on a grid
inpainting_grid = make_grid(all_inpaintings, nrow=4, pad_value=1.)
grid_as_np = inpainting_grid.permute(1, 2, 0).numpy()
# If NP returns out of range values for pixels, clip values
plt.imshow(np.clip(grid_as_np, 0, 1))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.09592188397472717, 'fraction_numerical': 0.018380241240666284, 'mean_word_length': 5.1646017699115045, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 22, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '144698', 'license': 'MIT', 'n_tokens_mistral': 1252, 'n_tokens_neox': 1176, 'n_words': 340, 'path': 'example-img.ipynb', 'repo_name': 'philippeitis/neural-processes', 'size': '3482'}