About output from the pretrained model

[im-aksh]

Hey Hi,
I am writing to you as a user of your paper's pretrained model, specifically the model described in the code you provided. First of all, I would like to express my appreciation for your work and the effort you have put into developing this model.

Recently, I have been utilizing your pretrained model for a specific task in my research project. While I acknowledge the potential and effectiveness of the model, I must inform you that I am not getting correct output.

I have attached my code and the output, please look into it and let me know if there are changes needed in the code

`from future import print_function
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
from PIL import Image, ImageOps
from torchvision.transforms.functional import to_pil_image
from models import InpaintingModel
import lpips
import os
from skimage.metrics import peak_signal_noise_ratio as compare_psnr
from skimage.metrics import structural_similarity as compare_ssim

loss_fn_alex = lpips.LPIPS(net='alex')

Training settings

parser = argparse.ArgumentParser(description='PyTorch Video Inpainting with Background Auxiliary')
parser.add_argument('--bs', type=int, default=256, help='training batch size')
parser.add_argument('--lr', type=float, default=0.001, help='Learning Rate. Default=0.001')
parser.add_argument('--cpu', default=False, action='store_true', help='Use CPU to test')
parser.add_argument('--threads', type=int, default=1, help='number of threads for data loader to use')
parser.add_argument('--seed', type=int, default=67454, help='random seed to use. Default=123')
parser.add_argument('--gpus', default=0, type=int, help='number of GPUs')
parser.add_argument('--threshold', type=float, default=0.8)
parser.add_argument('--img_path', type=str, default="D:/FYP/input_image/input.jpg")
parser.add_argument('--mask_path', type=str, default="D:/FYP/input_mask/00015.png")
parser.add_argument('--model', default='C:/FYP/RN-master/pretrained_model/x_admin.cluster.localRN-0.8RN-Net_bs_14_epoch_3.pth', help='pretrained base model')
parser.add_argument('--save', default=True, action='store_true', help='If save test images')
parser.add_argument('--save_path', type=str, default='C:/FYP/RN-master/output')
parser.add_argument('--input_size', type=int, default=512, help='input image size')
parser.add_argument('--l1_weight', type=float, default=1.0)
parser.add_argument('--gan_weight', type=float, default=.1)

opt = parser.parse_args()

def evaluate_single_image(image_path, mask_path, save=False, save_path=None):
# Load the model
device = torch.device('cpu' if opt.cpu else 'cuda')
model = InpaintingModel(g_lr=opt.lr, d_lr=(0.1 * opt.lr), l1_weight=opt.l1_weight, gan_weight=opt.gan_weight, iter=0, threshold=opt.threshold)
model.load_state_dict(torch.load(opt.model, map_location=device), strict=False)

pred, avg_lpips, mask, gt = eval_single_image(image_path, mask_path, model)

if save:
    image = Image.open(opt.img_path)
    mask = Image.open(opt.mask_path)
    inverted_mask = ImageOps.invert(mask)
    resized_mask = inverted_mask.resize(image.size, resample=Image.BILINEAR)
    masked_image = Image.composite(image, Image.new('RGB', image.size), resized_mask)
    masked_image.save(r'C:/FYP/RN-master/output/input.png')
    save_img(save_path, 'mask', mask)
    save_img(save_path, 'output', pred)
    save_img(save_path, 'gt', gt)

return avg_lpips

def eval_single_image(image_path, mask_path, model):
model.eval()
model.generator.eval()
avg_lpips = 0.

with torch.no_grad():
    gt = np.array(Image.open(image_path))
    mask = np.array(Image.open(mask_path))

    gt = torch.from_numpy(gt.transpose((2, 0, 1))).float().unsqueeze(0) / 255.0
    mask = torch.from_numpy(mask).unsqueeze(0).unsqueeze(0)  # Add extra dimensions for channel and batch

    # Resize the input tensor and mask tensor to match the expected input size
    gt = F.interpolate(gt, size=(opt.input_size, opt.input_size), mode='bilinear', align_corners=False)
    mask = F.interpolate(mask, size=(opt.input_size, opt.input_size), mode='nearest')

    gt, mask = Variable(gt), Variable(mask)

    prediction = model.generator(gt, mask)
    prediction = prediction * mask + gt * (1 - mask)
    avg_lpips = loss_fn_alex(prediction, gt).mean().item()

return prediction, avg_lpips, mask, gt

def save_img(path, name, img):
# img (H,W,C) or (H,W) np.uint8 or torch tensor
if isinstance(img, torch.Tensor):
img = to_pil_image(img.squeeze().cpu())
img.save(os.path.join(path, name + '.png'))

def main():
torch.manual_seed(opt.seed)

# Checking for GPU availability
use_cuda = torch.cuda.is_available()
device = torch.device("cuda" if use_cuda and not opt.cpu else "cpu")

# Evaluate single image
avg_lpips = evaluate_single_image(opt.img_path, opt.mask_path, save=opt.save, save_path=opt.save_path)

print("Average LPIPS: {:.4f}".format(avg_lpips))

if name == 'main':
main()
`