testing.py

import torch
from torchvision.transforms import ToPILImage
import torchvision.transforms as transforms
from PIL import Image
from networks import Generator


if __name__ == "__main__":

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    torch.backends.cuda.matmul.allow_tf32 = True
    NUM_WORKERS = 1
    img_size = 256
    n_block = 9
    subdir = 'maps'
    dom = 'A'
    opdom = 'B'
    f_num = 1029 #enter file number to test

    model_path = f'models/{subdir}/G_150_1_.pth'
    test_img_path = f'data/{subdir}/test/{f_num}_real_{dom}.png'


    transform = transforms.Compose([
        transforms.Resize((img_size, img_size)),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
    ])

    # load generator
    model = torch.compile(Generator(n_block=n_block)).to(device)
    model.load_state_dict(torch.load(
        model_path, map_location=device))
    model.eval()

    # load image
    
    input_img = Image.open(test_img_path)
    input = transform(input_img).unsqueeze(0).to(device)

    # generate output image
    with torch.no_grad():
        output = model(input)

    # save output image
    output_img = ToPILImage()(output.squeeze().cpu())
    output_img.save(f'outputs/c_{f_num}_{opdom}.jpg')