test.py

import time
from dataset import *
import os
import torch
from utils import *
import torch.nn as nn
import skimage.io
import argparse
import torch.distributed as dist
from torchvision import transforms
from PIL import Image
import tifffile
import pandas as pd
import matplotlib.pyplot as plt
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
from math import log10, sqrt
import torch.optim.lr_scheduler as lr_scheduler
from torch.distributions import normal
from cm2_model import *
from tensorboardX import SummaryWriter

parser = argparse.ArgumentParser(description='Train the network', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--dir_data', dest='dir_data')
parser.add_argument("--network", default='cm2net', help='multiwiener multifourier and cm2net')
parser.add_argument('--model_name', default='/cm2_test/', dest='model_name')
parser.add_argument('--reference_path')
parser.add_argument('--histogram', default='off')
parser.add_argument('--mask', default='off')
parser.add_argument('--fov', type=int,  default=6500)
parser.add_argument('--batch_size', type=int, default=1, dest='batch_size')
parser.add_argument("--local_rank", type=int, default=0, dest='local_rank')
parser.add_argument("--num_psf", type=int, default=9)
parser.add_argument("--ps", type=int, default=1)
parser.add_argument("--ks", type=float, default=10.0)
parser.add_argument("--epoch", type=int, default=103) # indicates the epoch of the best model so that it can be loaded
parser.add_argument('--dir_chck', default='./checkpoints/best_model_4_5e-4_150_BCEL2/', dest='dir_chck')
parser.add_argument("--distributed", type=bool, default=False, dest='distributed')
parser.add_argument('--lr', type=float, default=5e-5, dest='lr')
parser.add_argument('--mode', default='test', choices=['train', 'test'], dest='mode')


PARSER = Parser(parser)
args = PARSER.get_arguments()
PARSER.print_args()

torch.manual_seed(3407)
torch.cuda.empty_cache()
epoch = args.epoch

# make dir
dir_result_test = args.dir_data + args.model_name
if not os.path.exists(os.path.join(dir_result_test)):
    os.makedirs(os.path.join(dir_result_test))
# dir_result_gt = args.dir_data + '/gt/'
# if not os.path.exists(os.path.join(dir_result_gt)):
#     os.makedirs(os.path.join(dir_result_gt))

args.num_gpu = list(range(torch.cuda.device_count()))
torch.cuda.set_device(args.local_rank)
args.device=torch.device(f'cuda:{args.local_rank}')

# training data
if args.network == 'cm2net':
    # Create the complete dataset
    transform_train = transforms.Compose([ToTensor()])
    test_set = CM2Dataset_test(args.dir_data, transform=transform_train)
else:
    transform_train = transforms.Compose([Resize(), ToTensor()])
    test_set = MyDataset(args.dir_data, transform=transform_train)
length = len(test_set)
print('testing images:', length)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=args.batch_size, num_workers=0, shuffle=False)

## setup network TBD!
if args.network == 'multiwiener':
    psfs = skimage.io.imread(args.dir_data + '/psf_v11.tif')
    psfs = np.array(psfs)
    psfs = psfs.astype('float32') / psfs.max()
    psfs = psfs[:,57 * 2:3000, 94 * 2 + 156:4000 - 156]
    psfs = np.pad(psfs, ((0,0),(657, 657), (350, 350)))
    Ks = args.ks*np.ones((args.num_psf, 1, 1))
    deconvolution= MultiWienerDeconvolution2D(psfs,Ks).to(args.device)
    enhancement = RCAN(args.num_psf).to(args.device)
    model = LSVEnsemble2d(deconvolution, enhancement)

if args.network == 'multifourier':
    deconvolution = FourierDeconvolution2D_ds(args.num_psf,args.ps).to(args.device)
    enhancement = RCAN(args.num_psf).to(args.device)
    model = LSVEnsemble2d(deconvolution, enhancement)

if args.network == 'cm2net':
    layers = 20  # number of resblocks
    model = cm2net(numBlocks=layers, stackchannels=args.num_psf).to(args.device)  # the input is stack of 9 demixed views, output is one final reconstrution

#multiple gpu
if args.distributed:
    model = model.to(args.device)
    model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank], output_device=args.local_rank)
else:
    model = model.to(args.device)

## setup loss & optimization
ssim_loss = MS_SSIM(data_range=1, size_average=True, channel=1)
l1_loss = nn.L1Loss()
l2_loss = nn.MSELoss()
bce_loss = nn.BCELoss()
params = model.parameters()
optimizer = torch.optim.Adam(params, lr=args.lr)
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, 50, eta_min = 1e-6)


## load from checkpoints
model, losslogger = load(args.dir_chck, args.distributed, args.local_rank, model, optimizer, args.epoch, mode=args.mode)

ssim_test = []
psnr_test = []
speed_test = []
## validation phase
with torch.no_grad():
    model.eval()

    for batch, data in enumerate(test_loader, 1):
        # forward simulation(add noise)
        torch.cuda.current_stream().synchronize()
        t0 = time.time()
        meas = data['meas'].to(args.device)
        gt = data['gt'].to(args.device)
        if args.network == 'cm2net':
            demix_output, output = model(meas)
            torch.cuda.current_stream().synchronize()
            t1 = time.time()
        else:
            output = model(meas)
            torch.cuda.current_stream().synchronize()
            t1 = time.time()

        # loss = bce_loss(torch.squeeze(output, 1), gt)+l2_loss(torch.squeeze(output, 1), gt)
        output_n = (output - torch.min(output)) / (torch.max(output) - torch.min(output))
        gt_n = (gt - torch.min(gt)) / (torch.max(gt) - torch.min(gt))
        ssim = ssim_loss(output_n, gt_n.unsqueeze(1))
        psnr = 20 * torch.log10(torch.max(output) / sqrt(l2_loss(torch.squeeze(output, 1), gt)))

        # get losses
        ssim_test.append(ssim.item())
        psnr_test.append(psnr.item())
        speed_test.append(t1-t0)
        print('TEST: EPOCH %d: SAMPLE %04d/%04d: PSNR: %.4f SSIM: %.4f TIME: %.4f'
              % (epoch, batch, length, psnr.item(), ssim.item(), t1-t0))

        x_recon = output.data.cpu().numpy()
        im_recon = (np.clip(x_recon[-1, ...] / np.max(x_recon[-1, ...]), 0, 1) * 255).astype(np.uint8)
        tifffile.imwrite((dir_result_test + str(batch) + '_recon' + '.tif'), im_recon.squeeze())


    if args.network == 'multifourier':
        psfs_re = model.deconvolution.psfs_re.detach().cpu().numpy()
        psfs_im = model.deconvolution.psfs_im.detach().cpu().numpy()
        psf_freq = psfs_re + psfs_im * 1j
        psf = np.fft.ifftshift(np.fft.irfft2(psf_freq, axes=(-2, -1)))
        psf = (psf / np.abs(psf).max() * 65535.0).astype('int16')
        tifffile.imwrite((dir_result_test + str(epoch) + 'multifourier_psf' + '.tif'), psf, photometric='minisblack')
        psf_mip = np.max(psf, 0).squeeze()
        tifffile.imwrite((dir_result_test + str(epoch) + 'multifourier_psf_mip' + '.tif'), psf_mip, photometric='minisblack')

    if args.network == 'multiwiener':
        psfs = model.deconvolution.psfs.detach().cpu().numpy()
        psf = (psfs / np.abs(psfs).max() * 65535.0).astype('int16')
        tifffile.imwrite((dir_result_test  + 'multiwirner_psf' + '.tif'), psf, photometric='minisblack')
        psf_mip = np.max(psf, 0).squeeze()
        # psf_mip = (psf_mip / np.abs(psf_mip).max() * 65535.0).astype('int16')
        tifffile.imwrite((dir_result_test +  'multiwirner_psf_mip' + '.tif'), psf_mip, photometric='minisblack')
        # Create a dictionary with the recorded data

# Save records to CSV file
data = {'SSIM': ssim_test, 'PSNR': psnr_test, 'Speed':speed_test}
df = pd.DataFrame(data)
df.to_csv(dir_result_test + 'testing_records.csv', index=False)
print('TEST:AVERGARED: PSNR: %.4f SSIM: %.4f Speed: %.4f' % (np.mean(psnr_test), np.mean(ssim_test), np.mean(speed_test)))