GTCTV-DPC_msi.py

import os
import time
import glob
import torch 
import logging

import numpy as np
from scipy import io

import sys
sys.path.append("./utils/")
from utils_log import logger_info, logger_close
from utils_data import initial_seed
from utils_data import calculate_psnr_mdi, calculate_ssim_mdi, imgssave
from utils_lowrank import tensor_svd, shrinkage
from utils_lowrank import dct_unit, idct_unit
from utils_lowrank import diff_mdi, diff_mdi_T, KKT_mdi

device = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")

def svt_lrtc(tensor, tau, 
             unitary_operation=dct_unit, 
             inv_unitary_operation=idct_unit,
             ):
    u, s, v = tensor_svd(unitary_operation(tensor))
    ss = shrinkage(s, tau, mode="soft")
    try:
        idx = torch.max(torch.where(ss > 0)[-1]) + 1
        # if the unitary_operation is DFT, the dtype of u, v is complexfloat32, while the dtype of ss is float32.
        return inv_unitary_operation(u[..., :idx] @ torch.diag_embed(ss[..., :idx]).to(u.dtype) @ v[..., :idx, :])
    except:
        # when the shrinkage threshold is too large, the result is all zeros
        # if the shrinkage result is all zeros, we return a zero tensor with the same shape as the input tensor
        return torch.zeros_like(tensor)

def deep_pnp_solver(tensor, net, sigma):
    # tensor: (T, C, H, W)
    T, _, H, W = tensor.shape
    noise_level_map = torch.ones(T, 1, H, W).mul_(sigma).float().to(tensor.device)
    input = torch.cat([tensor, noise_level_map], dim=1)
    with torch.no_grad():
        return net(input)

def fidelity_prox(tensor, Y, mask):
    # proximal operator for completion problem
    return (1 - mask) * tensor + Y

def dys_restore(
        observe_tensor, mask, # fidelity term
        tau, beta, rho, orders, inner_iter, # tctv term
        model, sigma_init, sigma_min, alpha, # deep denoiser
        ori_tensor=None, save_dir=None, logger=None, 
        check_iter=1, max_iter=500, tole=1e-4,
        ):
    
    # observe_tensor: (T, C, H, W), range: [0, 255]
    # change the range to [0, 1]
    Y = observe_tensor / 255
    initial_seed(1000)

    # variable initialization
    z = torch.randn_like(Y)
    z = (1 - mask) * z + Y
    out_tensor = z.clone()

    # here sigma is the noise level, which divided by 255
    sigma = sigma_init

    # variable initialization for proximal solution in Low-rank with dim-1 TV
    n = len(orders)
    Gs = [diff_mdi(z, axis=order) for order in orders]
    Bs = [torch.zeros_like(z) for _ in orders]

    # prepare the KKT tensor for the total variation
    # only use the shape of the tensor to compute the KKT tensor
    KKT_fft = sum(KKT_mdi(z, axis=order) for order in orders)
    
    # variable initialization for proximal solution in fidelity term
    used_time = 0
    psnrs = []
    ssims = []

    # iterative optimization
    for i in range(max_iter):
        # compute the PSNR, SSIM
        if ori_tensor is not None:
            if i % check_iter == 0:
                # (T, C, H, W) -> (H, W, C, T)
                temp_out = out_tensor.cpu().permute(2, 3, 1, 0).numpy() * 255
                temp_ori = ori_tensor.cpu().permute(2, 3, 1, 0).numpy()

                psnr_value = calculate_psnr_mdi(temp_out, temp_ori)
                ssim_value = calculate_ssim_mdi(temp_out, temp_ori)
                psnrs.append([i, psnr_value])
                ssims.append([i, ssim_value])

                logger.info("Iter: {:03d}, PSNR: {:.3f}, SSIM: {:.4f}".format(i, psnr_value, ssim_value))
            
                if save_dir is not None:
                    print("Saving the MSI...")
                    save_dir_i = os.path.join(save_dir, f"iter_{i}_psnr_{psnr_value:.3f}_ssim_{ssim_value:.4f}")
                    imgssave(temp_out, save_dir_i, metric=True, ori_imgs=temp_ori)
                    print("MSI saved.")
        
        # record the time and update the input sigma value
        start_time = time.time()

        # DYS step
        # proximal operator for Low-rank with dim-{orders} TV
        for inner_it in range(inner_iter):
            ### tctv
            # update x_B
            H = sum(diff_mdi_T(beta * Gs[index] - Bs[index], axis=orders[index]) for index in range(n))
            Nominator = torch.fft.fftn(tau * H + z)
            Denominator = tau * beta * KKT_fft + 1
            x_B = torch.real(torch.fft.ifftn(Nominator / Denominator))

            # update the gradient tensors and lagrange multipliers
            for index in range(n):
                temp = diff_mdi(x_B, axis=orders[index]) + Bs[index] / beta
                Gs[index] = svt_lrtc(temp, 1 / (n * beta))
                Bs[index] += beta * (diff_mdi(x_B, axis=orders[index]) - Gs[index])
            
            beta = min(beta * rho, 1e10)
            if inner_it > 0:
                inner_tol = torch.linalg.norm(x_B - x_B_old) / torch.linalg.norm(x_B_old)
                if inner_tol < tole:
                    logger.info("Inner iteration converged at iter: {}".format(inner_it+1))
                    break
            
            x_B_old = x_B.clone()

        x_C = deep_pnp_solver(x_B, model, sigma)
        temp_x_A = 2 * x_B - z - tau * alpha * (x_B - x_C) # 2 * J_{tau * B} - I - tau * C(J_{tau * B})
        # tau = max(tau / rho, 1e-4)
        x_A = fidelity_prox(temp_x_A, Y, mask) # A(x_B) fidelity term of completion problem

        sigma = max(sigma / rho, sigma_min)
        del x_C, temp_x_A

        # update the z
        # z += lamb_t * (x_A - x_B)
        if i < 100:
            z += (x_A - x_B) # set lamb_t to 1.
        else:
            z += 100 / i * (x_A - x_B) # set lamb_t to 1 / t.
        
        # record the time
        used_time += time.time() - start_time

        tol = torch.linalg.norm(out_tensor - x_A) / torch.linalg.norm(out_tensor)
        if tol < tole:
            break
        out_tensor = x_A.clone()
    
    # compute the PSNR, SSIM
    if ori_tensor is not None:
        # (T, C, H, W) -> (T, H, W, C)
        temp_out = out_tensor.cpu().permute(2, 3, 1, 0).numpy() * 255
        temp_ori = ori_tensor.cpu().permute(2, 3, 1, 0).numpy()

        psnr_value = calculate_psnr_mdi(temp_out, temp_ori)
        ssim_value = calculate_ssim_mdi(temp_out, temp_ori)

        psnrs.append([i+1, psnr_value])
        ssims.append([i+1, ssim_value])

        logger.info("Iter: {:03d}, PSNR: {:.3f}, SSIM: {:.4f}".format(i+1, psnr_value, ssim_value))
        logger.info("Total time: {:.5f}s".format(used_time))
        if save_dir is not None:
            print("Saving the MSI...")
            # save the last result
            save_dir_i = os.path.join(save_dir, f"iter_{i}_psnr_{psnr_value:.3f}_ssim_{ssim_value:.4f}")
            imgssave(temp_out, save_dir_i, metric=True, ori_imgs=temp_ori)
        
    return out_tensor, psnrs, ssims, used_time

param_file = "./params/SPC09_gray.pth"
from utils_drunet.network_unet import UNetRes as net
model = net(
    in_nc=2, 
    out_nc=1, 
    nc=[64, 128, 256, 512], 
    nb=4, 
    act_mode="R", 
    downsample_mode="strideconv", 
    upsample_mode="convtranspose", 
    bias=False,
    interpolation_mode="bilinear"
)
state_dict = torch.load(param_file, weights_only=True)
model.load_state_dict(state_dict, strict=True)
model.eval()

data_root = './test_datasets/msis'
data_files = glob.glob(os.path.join(data_root, "*.mat"))

sampling_rates = [0.05, 0.1, 0.2]
beta, rho, sigma_init, alpha = 1e-4, 1.02, 0.05, 1.0
inner_iter = 8
sigma_min = 1e-3
tau = 1.

res_dir_root = "./exp_results/msis"
for sampling_rate in sampling_rates:
    # create the log and result directories
    temp_dir_root = res_dir_root + f'_{sampling_rate:.2f}'
    log_dir = os.path.join(temp_dir_root, "logs")
    if not os.path.exists(log_dir):
        os.makedirs(log_dir)
    res_dir = os.path.join(temp_dir_root, "results")
    if not os.path.exists(res_dir):
        os.makedirs(res_dir)
    
    psnrs_sum = 0
    ssims_sum = 0
    for data_file in data_files:
        img_name = os.path.basename(data_file).split("/")[-1].split(".")[0]
        # load the data
        dense_tensor = io.loadmat(data_file)['dense_tensor']
        binary_tensor = io.loadmat(data_file)[f'mask_sr_{sampling_rate:.2f}']
        sparse_tensor = io.loadmat(data_file)[f'sparse_tensor_sr_{sampling_rate:.2f}']
        # convert the (256, 256, 31) array to (31, 1, 256, 256) tensor
        # range [0, 255]
        dense_tensor = torch.from_numpy(dense_tensor).permute(2, 0, 1).unsqueeze(1).float()
        binary_tensor = torch.from_numpy(binary_tensor).permute(2, 0, 1).unsqueeze(1).float()
        sparse_tensor = torch.from_numpy(sparse_tensor).permute(2, 0, 1).unsqueeze(1).float()

        logger_info(log_dir, log_path=os.path.join(log_dir, f"{img_name}.log"))
        result_logger = logging.getLogger(log_dir)
        result_logger.info("--------------------------------------------------------------------------------------------------------")
                
        save_dir_name = os.path.join(res_dir, f"{img_name}")
        if not os.path.exists(save_dir_name):
            os.makedirs(save_dir_name)

        out_tensor, psnrs, ssims, used_time = dys_restore(
            sparse_tensor.to(device), mask=binary_tensor.to(device), 
            tau=tau, beta=beta, rho=rho, orders=[0, 2, 3], inner_iter=inner_iter,
            model=model.to(device), sigma_init=sigma_init, sigma_min=sigma_min, alpha=alpha, 
            ori_tensor=dense_tensor, logger=result_logger, save_dir=save_dir_name, 
            check_iter=100, max_iter=200, tole=1e-4,
            )
        
        psnrs_sum += psnrs[-1][-1]
        ssims_sum += ssims[-1][-1]
        # save the results
        res_save_dir_name = os.path.join(save_dir_name, 'res')
        if not os.path.exists(res_save_dir_name):
            os.makedirs(res_save_dir_name)
        # convert the (31, 1, 256, 256) tensor to (256, 256, 1, 31) array with range [0, 255] of np.uint8
        save_img = out_tensor.cpu().detach().permute(2, 3, 1, 0).numpy() * 255
        save_img = np.uint8(save_img.clip(0, 255).round())
        np.save(os.path.join(res_save_dir_name, "output.npy"), save_img)
        np.save(os.path.join(res_save_dir_name, "psnrs.npy"), psnrs)
        np.save(os.path.join(res_save_dir_name, "ssims.npy"), ssims)
        result_logger.info("--------------------------------------------------------------------------------------------------------")
        logger_close(result_logger)
    
    # save the average results
    avg_psnr = psnrs_sum / len(data_files)
    avg_ssim = ssims_sum / len(data_files)
    logger_info(log_dir, log_path=os.path.join(log_dir, f"avg-psnr_{avg_psnr:.3f}_ssim_{avg_ssim:.4f}.log"))
    temp_logger = logging.getLogger(log_dir)
    temp_logger.info("--------------------------------------------------------------------------------------------------------")
    temp_logger.info("Average PSNR: {:.3f}, Average SSIM: {:.4f}".format(avg_psnr, avg_ssim))
    temp_logger.info("--------------------------------------------------------------------------------------------------------")
    logger_close(temp_logger)