# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'

from PIL import Image
import numpy as np
import torch
from torch.utils.data import DataLoader
from input_pipeline import PairDataset
# -

# # Create an input pipeline

# +
dataset = PairDataset(
    first_dir='',
    second_dir='',
    num_samples=1000,
    image_size=64
)

data_loader = DataLoader(
    dataset=dataset,
    batch_size=30, shuffle=True,
    num_workers=1, pin_memory=True
)
# -

# # Get random images

# +
for i, (x, y) in enumerate(data_loader):
    break
    
x = 255.0*x.permute(0, 2, 3, 1).numpy()
y = 255.0*y.permute(0, 2, 3, 1).numpy()
# -

# # Show a grid

images1 = np.concatenate(x, axis=0)
images2 = np.concatenate(y, axis=0)
Image.fromarray(np.concatenate([images1, images2], axis=1).astype('uint8'))

StarcoderdataJupyterScriptsDedupFiltered