<jupyter_start><jupyter_text>Libraries<jupyter_code>import time
from utils.data_load import DataLoad
from models.model import create_model
import torch
import os
import torchvision
from torch.utils import data
import torchvision.transforms as transforms

from piqa import psnr, ssim 
from piqa.utils.functional import gaussian_kernel<jupyter_output><empty_output><jupyter_text>Settings<jupyter_code>class Opion():
    
    def __init__(self):
            
        self.dataroot = r'/path/to/images/'  # image dataroot
        self.maskroot = r'/path/to/masks/'  # mask dataroot
        self.batchSize = 1   # Need to be set to 1
        self.fineSize = 256  # image size
        self.input_nc = 3  # input channel size for first stage
        self.input_nc_g = 6  # input channel size for second stage
        self.output_nc = 3  # output channel size
        self.ngf = 64  # inner channel
        self.ndf = 64  # inner channel
        self.which_model_netD = 'basic'  # patch discriminator
        self.which_model_netF = 'feature'  # feature patch discriminator
        self.which_model_netG = 'unet_csa'  # seconde stage network
        self.which_model_netP = 'unet_256'  # first stage network
        self.triple_weight = 1
        self.name = 'CSA_inpainting'
        self.n_layers_D = '3'  # network depth
        self.gpu_ids = [0]
        self.model = 'csa_net'
        self.checkpoints_dir = r'/path/to/checkpoints/'
        self.norm = 'instance'
        self.fixed_mask = 1
        self.use_dropout = False
        self.init_type = 'normal'
        self.mask_type = 'random'
        self.lambda_A = 100
        self.threshold = 5/16.0
        self.stride = 1
        self.shift_sz = 1  # size of feature patch
        self.mask_thred = 1
        self.bottleneck = 512
        self.gp_lambda = 10.0
        self.ncritic = 5
        self.constrain = 'MSE'
        self.strength = 1
        self.init_gain = 0.02
        self.cosis = 1
        self.gan_type = 'lsgan'
        self.gan_weight = 0.2
        self.ssim_weight = 100
        self.lorentzian_weight = 10
        self.overlap = 4
        self.skip = 0
        self.display_freq = 10
        self.print_freq = 2
        self.save_latest_freq = 5
        self.save_epoch_freq = 2
        self.continue_train = False
        self.epoch_count = 1
        self.phase = 'test'  # or train
        self.which_epoch = '120'
        self.niter = 2
        self.niter_decay = 4
        self.beta1 = 0.5
        self.lr = 0.0002
        self.lr_policy = 'lambda'
        self.lr_decay_iters = 50
        self.isTrain = False
        self.ssim_loss = True  # True or False if we want to use / don't use SSIM loss additionally
        self.lorentzian_loss = False  # True or False if we want to use / don't use Lorentzian loss additionally
        self.l1_weight = 0.1<jupyter_output><empty_output><jupyter_text>Test Dataset<jupyter_code>opt = Opion()
transform_mask = transforms.Compose(
    [transforms.Resize((opt.fineSize, opt.fineSize)),
     transforms.ToTensor(),
    ])
transform = transforms.Compose(
    [
     transforms.Resize((opt.fineSize, opt.fineSize)),
     transforms.ToTensor(),
     transforms.Normalize(mean=[0.5] * 3, std=[0.5] * 3)])

dataset_test = DataLoad(opt.dataroot, opt.maskroot, transform, transform_mask)
iterator_test = (data.DataLoader(dataset_test, batch_size=opt.batchSize, shuffle=True))
print(len(dataset_test))
model = create_model(opt)
total_steps = 0<jupyter_output>100
csa_net
initialize network with normal<jupyter_text>MetricsMetrics that were used to measure the correctness of the algorithm:PSNR - (Peak Signal-to-Noise Ratio)SSIM - (Structure Similarity Index Measure)MSE - (Mean Square Error)<jupyter_code>criterionMSE = torch.nn.MSELoss()<jupyter_output><empty_output><jupyter_text>Model Loading<jupyter_code>load_epoch = 120
model.load(load_epoch)<jupyter_output><empty_output><jupyter_text>Model Testing<jupyter_code>psnr_results = []
ssim_results = []
mse_results = []
save_dir = '/path/to/save/'
if os.path.exists(save_dir) is False:
    os.makedirs(save_dir)

for count, (image, mask) in enumerate(iterator_test):
    iter_start_time = time.time()
    image=image.cuda()
    mask=mask.cuda()
    mask=mask[0][0]
    mask=torch.unsqueeze(mask,0)
    mask=torch.unsqueeze(mask,1)
    mask=mask.bool()

    model.set_input(image,mask)
    model.set_gt_latent()
    model.test()
    real_A,real_B,fake_B=model.get_current_visuals()
    pic = (torch.cat([real_A, real_B,fake_B], dim=0) + 1) / 2.0
    torchvision.utils.save_image(pic, '%s/Image_(%d)_(%dof%d).jpg' % (
    save_dir, count, count + 1, len(dataset_test)), nrow=1)
    
    # Calculate MSE
    acc_mse = criterionMSE(real_B, fake_B).item()
    print("For {} batch:".format(count + 1)) 
    print("MSE ==> {}".format(acc_mse))
    mse_results.append(acc_mse)
    # Calculate PSNR
    acc_psnr = psnr.psnr(real_B, fake_B).item()
    print("{} ==> {}".format("PSNR", acc_psnr))
    psnr_results.append(acc_psnr)
    # Calculate SSIM
    kernel = gaussian_kernel(11).repeat(3, 1, 1).cuda()
    acc_ssim = ssim.ssim(real_B, fake_B, kernel)
    print("{} ==> {}".format("SSIM", acc_ssim[0].item()))
    ssim_results.append(acc_ssim[0].item())

# MSE result for test data
mse_test_data = sum(mse_results) / len(mse_results)
print("MSE for test data is: {}".format(mse_test_data))

# PSNR result for test data
psnr_test_data = sum(psnr_results) / len(psnr_results)
print("PSNR for test data is: {}".format(psnr_test_data))

# SSIM result for test data
ssim_test_data = sum(ssim_results) / len(ssim_results)
print("SSIM for test data is: {}".format(ssim_test_data))<jupyter_output>For 1 batch:
MSE ==> 0.053971774876117706
PSNR ==> 12.678332328796387
SSIM ==> 0.4671025574207306
For 2 batch:
MSE ==> 0.048840295523405075
PSNR ==> 13.112215995788574
SSIM ==> 0.4886021912097931
For 3 batch:
MSE ==> 0.04842561483383179
PSNR ==> 13.149248123168945
SSIM ==> 0.49265050888061523
For 4 batch:
MSE ==> 0.04408326745033264
PSNR ==> 13.55726146697998
SSIM ==> 0.5383450984954834
For 5 batch:
MSE ==> 0.04372094199061394
PSNR ==> 13.593103408813477
SSIM ==> 0.5345590114593506
For 6 batch:
MSE ==> 0.03730323165655136
PSNR ==> 14.282533645629883
SSIM ==> 0.651111900806427
For 7 batch:
MSE ==> 0.123477503657341
PSNR ==> 9.084121704101562
SSIM ==> 0.38831886649131775
For 8 batch:
MSE ==> 0.06932743638753891
PSNR ==> 11.590948104858398
SSIM ==> 0.5666568875312805
For 9 batch:
MSE ==> 0.027314377948641777
PSNR ==> 15.636085510253906
SSIM ==> 0.5340876579284668
For 10 batch:
MSE ==> 0.055225010961294174
PSNR ==> 12.578640937805176
SSIM ==> 0.5490313768386841
For 11 batch:
MSE ==> 0.0358[...]
{'fraction_non_alphanumeric': 0.11809839025124116, 'fraction_numerical': 0.09688581314878893, 'mean_word_length': 3.3168831168831168, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 57, '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': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '415874', 'license': 'MIT', 'n_tokens_mistral': 2869, 'n_tokens_neox': 2459, 'n_words': 503, 'path': 'test.ipynb', 'repo_name': 'klaudiaplk/ImageInpainting', 'size': '6647'}