FNorm4RGB.py

import os
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
import torch
from torch import nn, optim 
import torch.nn.functional as F
import numpy as np 
import matplotlib.pyplot as plt 
from skimage.metrics import peak_signal_noise_ratio, normalized_root_mse, structural_similarity
import rff
import random
import copy
from tqdm import tqdm
from PIL import Image
import argparse
from scipy.ndimage import gaussian_filter1d
dtype = torch.cuda.FloatTensor

def set_random_seed(seed):
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    np.random.seed(seed)
    random.seed(seed)


class MLPLayer(nn.Module):
    def __init__(self, in_features, out_features, bias=True, is_first=False, omega_0=1.0*np.pi):
        super().__init__()
        self.omega_0 = omega_0
        self.is_first = is_first
        self.in_features = in_features
        self.linear = nn.Linear(in_features, out_features, bias=bias)
        self.init_weights()
    
    def init_weights(self):
        with torch.no_grad():
            if self.is_first:
                nn.init.uniform_(self.linear.weight, -1 / self.in_features, 1 / self.in_features)
            else:
                # self.linear.weight.uniform_(-np.sqrt(6 / self.in_features) / self.omega_0, 
                #                              np.sqrt(6 / self.in_features) / self.omega_0)
                nn.init.xavier_uniform_(self.linear.weight, gain=nn.init.calculate_gain('relu'))
                # nn.init.xavier_uniform_(self.linear.weight, gain=nn.init.calculate_gain('leaky_relu', param=0.1))

    def forward(self, input):
        output = self.linear(input)
        output = F.relu(output) #tanh、hardtanh、softplus、relu、sin
        # output = F.leaky_relu(output, negative_slope=0.1)
        return output

mid_channel = 512
rank = 256
posdim = 128
class Network(nn.Module):
    def __init__(self, rank, mid_channel, posdim):
        super(Network, self).__init__()

        U_net = nn.Sequential(MLPLayer(posdim, mid_channel, is_first=True),
                              MLPLayer(mid_channel, mid_channel, is_first=False),
                              MLPLayer(mid_channel, mid_channel, is_first=False),
                              nn.Linear(mid_channel, rank))
        self.U_Nets = nn.ModuleList([copy.deepcopy(U_net) for _ in range(3)])

        V_net = nn.Sequential(MLPLayer(posdim, mid_channel, is_first=True),
                              MLPLayer(mid_channel, mid_channel, is_first=False),
                              MLPLayer(mid_channel, mid_channel, is_first=False),
                              nn.Linear(mid_channel, rank))
        self.V_Nets = nn.ModuleList([copy.deepcopy(V_net) for _ in range(3)])
        
        self.encoding = rff.layers.GaussianEncoding(alpha=1.0, sigma=8.0, input_size=1, encoded_size=posdim//2) #[sigma=5,8]
    
    def forward(self, U_input, V_input):
        outputs = []
        outputs_U = []
        outputs_V = []
        for i in range(3):
            U = self.U_Nets[i](self.encoding(self.normalize_to_01(U_input)))
            V = self.V_Nets[i](self.encoding(self.normalize_to_01(V_input)))
            channelImg = U @ V.t()
            outputs.append(channelImg)
            outputs_U.append(U)
            outputs_V.append(V)
        output = torch.stack(outputs, dim=-1)
        Us = torch.stack(outputs_U, dim=-1)
        Vs = torch.stack(outputs_V, dim=-1)
        return output, Us, Vs
    
    def normalize_to_01(self, tensor):
        min_val = tensor.min()
        max_val = tensor.max()
        if max_val == min_val:
            return torch.zeros_like(tensor)
        normalized_tensor = (tensor - min_val) / (max_val - min_val)
        return normalized_tensor

def generate_random_mask(H, W, visible_ratio=0.1):
    num_visible = int(H * W * visible_ratio)
    all_positions = np.arange(H * W)
    visible_positions = np.random.choice(all_positions, size=num_visible, replace=False)
    mask = np.zeros(H * W, dtype=np.uint8)
    mask[visible_positions] = 1
    mask = mask.reshape(H, W, 1)
    mask = np.repeat(mask, 3, axis=2)
    mask = torch.from_numpy(mask).type(dtype).cuda()
    return mask


if __name__ == '__main__':
    set_random_seed(42)
    max_iter =  5001
    Schatten_q = 1.0
    image_names = ['Plane', 'Peppers', 'House', 'Sailboat']
    average_metrics = [0.0, 0.0, 0.0]
    OB_metrics = [0.0, 0.0, 0.0]
    loss_rec_dict = {'Plane': [], 'Peppers': [], 'House': [], 'Sailboat': []}
    color_dict = {'Plane': "#ef7262", 'Peppers': '#3498db', 'House': '#2ecc71', 'Sailboat': '#9b59b6'}

    parser = argparse.ArgumentParser(description="Process some integers.")
    parser.add_argument('--visible_ratio', type=float, required=True, help='The visible ratio parameter.')
    parser.add_argument('--save', action='store_true')
    args = parser.parse_args()

    for name in image_names:
        image_path = 'data/misc/'+name+'.tiff'
        image_gt = np.array(Image.open(image_path)).astype(np.float32) / 255.0        
        H, W, C = image_gt.shape
        mask = generate_random_mask(H, W, visible_ratio=args.visible_ratio)
        X = torch.from_numpy(image_gt).type(dtype).cuda()
        
        RGB_incomplete = (X*mask).cpu().detach().numpy()
        ob_psnr = peak_signal_noise_ratio(image_gt, RGB_incomplete, data_range=1.0)
        ob_ssim = structural_similarity(image_gt, RGB_incomplete, data_range=1.0, channel_axis=2)
        ob_nrmse = normalized_root_mse(image_gt, RGB_incomplete)
        print('SR:',args.visible_ratio, 'name:',name,
              'OB_PSNR: {:.3f}, OB_SSIM: {:.3f}, OB_NRMSE: {:.3f}'.format(ob_psnr, ob_ssim, ob_nrmse))
        OB_metrics = list(map(lambda x, y: x + y, OB_metrics, [ob_psnr, ob_ssim, ob_nrmse]))

        U_input = torch.from_numpy(np.array(range(1,H+1))).reshape(H,1).type(dtype)
        V_input = torch.from_numpy(np.array(range(1,W+1))).reshape(W,1).type(dtype)

        model = Network(rank, mid_channel, posdim).type(dtype)
        optimizer = optim.Adam([{'params': model.parameters(), 'weight_decay': 0.005}], #[0.005]
                                lr=0.001)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=max_iter, eta_min=0)
        
        best_metric = [0.0, 0.0, 0.0]
        with tqdm(total=max_iter) as pbar:
            for iter in range(max_iter):
                X_Out, Us, Vs = model(U_input, V_input)
                loss_rec = torch.norm(X_Out*mask - X*mask, p='fro', dim=(0, 1)).mean()
                
                U_input_eps = torch.normal(mean=U_input, std=1.0*torch.ones_like(U_input)) #std=1.0
                V_input_eps = torch.normal(mean=V_input, std=1.0*torch.ones_like(V_input))
                X_Out_eps, *_ = model(U_input_eps, V_input_eps)
                loss_eps = torch.norm(X_Out.detach()-X_Out_eps, p='fro', dim=(0, 1)).mean()

                # loss_rank = torch.norm(Us, p='fro', dim=(0, 1)).mean() + \
                #             torch.norm(Vs, p='fro', dim=(0, 1)).mean()
                loss_rank = torch.norm(Us, p=2, dim=0).pow(Schatten_q).sum() +\
                            torch.norm(Vs, p=2, dim=0).pow(Schatten_q).sum()
                
                loss = 1.0*loss_rec + 0.01*loss_eps + 0.001*loss_rank #[1.0, 0.01 0.001]
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()
                scheduler.step()
                pbar.set_postfix({'loss_rec': f"{loss_rec.item():.4f}", 
                                'loss_eps': f"{loss_eps.item():.4f}", 
                                'loss_rank': f"{loss_rank.item():.4f}"})
                pbar.update()
                loss_rec_dict[name].append(loss_rec.item())
                continue
                if iter % 500 == 0 and iter != 0:
                    psnr = peak_signal_noise_ratio(image_gt, X_Out.cpu().detach().numpy(), data_range=1.0)
                    ssim = structural_similarity(image_gt, X_Out.cpu().detach().numpy(), data_range=1.0, channel_axis=2)
                    nrmse = normalized_root_mse(image_gt, X_Out.cpu().detach().numpy())
                    if psnr >= best_metric[0]:
                        print('SR:',args.visible_ratio, 'name:',name, 'iteration:', iter, 'PSNR', 
                              psnr, 'SSIM', ssim, 'NRMSE', nrmse)
                        best_metric = [psnr, ssim, nrmse]
                        if args.save:
                            output_folder = os.path.join('output/Ours/inpainting/Image')
                            if not os.path.exists(output_folder):
                                os.makedirs(output_folder)
                            output_path = os.path.join(output_folder, name + f'_SR{args.visible_ratio:.2f}_psnr{psnr:.3f}_inpainting.png')
                            img = Image.fromarray((np.clip(X_Out.cpu().detach().numpy(),0,1) * 255).astype(np.uint8))
                            img.save(output_path)

                        save_path = f'model_weights_{name}.pth'
                        torch.save(model.state_dict(), save_path)
                        print(f"model saved in {save_path}")
                    continue

                    plt.figure(figsize=(15,45))
                    plt.subplot(131)
                    plt.imshow(np.clip((X*mask).cpu().detach().numpy(),0,1))
                    plt.title('in')

                    plt.subplot(132)
                    plt.imshow(image_gt)
                    plt.title('gt')
            
                    plt.subplot(133)
                    plt.imshow(np.clip(X_Out.cpu().detach().numpy(),0,1))
                    plt.title('out')
                    plt.show()

        average_metrics = list(map(lambda x, y: x + y, average_metrics, best_metric))

    print(f'SR:{args.visible_ratio:.2f}', 
          'OB_PSNR: {}, OB_SSIM: {}, OB_NRMSE: {}'.format(*['{:.3f}'.format(metric / len(image_names)) 
                                                   for metric in OB_metrics]))
    print(f'SR:{args.visible_ratio:.2f}', 
          'PSNR: {}, SSIM: {}, NRMSE: {}'.format(*['{:.3f}'.format(metric / len(image_names)) 
                                                   for metric in average_metrics]))
    
    import pickle
    with open('Err_RGB.pkl', 'wb') as f:
        pickle.dump(loss_rec_dict, f, protocol=pickle.HIGHEST_PROTOCOL)
    with open('Err_RGB.pkl', 'rb') as f:
        loss_rec_dict = pickle.load(f)

    # 绘制折线图
    fontsize = 20
    plt.rcParams.update({'font.size': fontsize})  # 设置全局字体大小
    plt.figure(figsize=(6, 5))
    for dataName in color_dict.keys():
        # loss_rec_dict[dataName] = gaussian_filter1d(loss_rec_dict[dataName], sigma=2)
        plt.plot(range(1, len(loss_rec_dict[dataName])+1), loss_rec_dict[dataName], 
                linestyle='-', color=color_dict[dataName], alpha=0.9, label=f'{dataName}')

    # 添加标题和标签
    plt.title('RGB', fontsize=fontsize)
    plt.xlabel('Iteration', fontsize=fontsize)
    plt.ylabel('F-Norm', fontsize=fontsize)

    # 添加网格和图例
    plt.grid(True, linestyle='--', alpha=0.6)
    plt.legend()

    # 显示图形
    plt.tight_layout()
    plt.savefig(f'Err_RGB.png')