train_settings/dvd/evaluation.py

import os

import numpy as np
import torch as th
import torch.nn.functional as F
from torch.utils.data import DataLoader
from tqdm import tqdm


from train_settings.dvd.feature_backbones.VGG_features import VGGPyramid

from utils_flow.visualization_utils import visualize, visualize_dewarping

from .eval_utils import extract_raw_features_single, extract_raw_features_single2
from .improved_diffusion import dist_util
from .improved_diffusion.gaussian_diffusion import GaussianDiffusion
import torch
from torchvision.utils import save_image as tv_save_image

def prepare_data(settings, batch_preprocessing, SIZE, data):
    if 'source_image_ori' in data:
        source_vis = data['source_image_ori']  # B, C, 512, 512 torch.uint8 cpu
    else:
        source_vis = data['source_image']
    if 'target_image' in data:
        target_vis = data['target_image']
    else:
        target_vis = None
    _, _, H_ori, W_ori = source_vis.shape

    # data = batch_preprocessing(data)

    source = data['source_image'].to(dist_util.dev())  # [1, 3, 914, 1380]  torch.float32
    if 'source_image_0' in data:
        source_0 = data['source_image_0'].to(dist_util.dev())
    else:
        source_0 = None
    if 'target_image' in data:  
        target = data['target_image']  # [1, 3, 914, 1380]  torch.float32
    else:
        target = None 
    if 'flow_map' in data:    
        batch_ori = data['flow_map']   # [1, 2, 914, 1380]  torch.float32
    else:
        batch_ori = None    
    if 'flow_map_inter' in data:    
        batch_ori_inter = data['flow_map_inter']   # [1, 2, 914, 1380]  torch.float32
    else:
        batch_ori_inter = None    
    if target is not None:
        target = F.interpolate(target, size=512, mode='bilinear', align_corners=False) # [1, 3, 512, 512] 
        target_256 = data['target_image_256'].to(dist_util.dev()) # [1, 3, 256, 256]
    else:
        target = None
        target_256 = None
    # source = F.interpolate(source, size=512, mode='bilinear', align_corners=False) #[1, 3, 512, 512] 
    # source_256 = data['source_image_256'].to(dist_util.dev()) # [1, 3, 256, 256]
    
    if settings.env.eval_dataset == 'hp-240':# false
        source_256 = source
        target_256 = target

    else: # true
        data['source_image_256'] = torch.nn.functional.interpolate(input=source.float(), size=(256, 256), mode='area')
        source_256 = data['source_image_256'].to(dist_util.dev())
        
        if 'target_image_256' in data:
            target_256 = data['target_image_256']
        else:
            target_256 = None
    if 'correspondence_mask' in data:
        mask = data['correspondence_mask'] # torch.bool [1, 914, 1380]
    else:
        mask = torch.ones((1, 512, 512), dtype=torch.bool).to(dist_util.dev()) # None

    return data, H_ori, W_ori, source, target, batch_ori, batch_ori_inter, source_256, target_256, source_vis, target_vis, mask, source_0



def run_sample_lr_dewarping(
    settings, logger, diffusion, model, radius, source, feature_size, 
    raw_corr, init_flow, c20, source_64, pyramid, doc_mask, 
    seg_map_all=None, textline_map=None, init_feat=None
):
    # init_flow = init_flow * feature_size
    # coords = initialize_flow(init_flow.shape[0], feature_size, feature_size, dist_util.dev())
    # coords_warped = coords + init_flow

    # local_corr = local_Corr(
    #     raw_corr.view(1, 1, feature_size, feature_size, feature_size, feature_size).to(dist_util.dev()),
    #     coords_warped.to(dist_util.dev()),
    #     radius,
    # )

    # local_corr = F.interpolate(
    #     local_corr.view(1, (2 * radius + 1) ** 2, feature_size, feature_size),
    #     size=feature_size,
    #     mode='bilinear',
    #     align_corners=True,
    # )

    # init_flow = F.interpolate(init_flow, size=feature_size, mode='bilinear', align_corners=True)
    # init_flow /= feature_size

    
    model_kwsettings = {'init_flow': init_flow, 'src_feat': c20, 'src_64':None, 
                        'y512':source, 'tmode':settings.env.train_mode,
                        'mask_cat': doc_mask,
                        'init_feat': init_feat,
                        'iter': settings.env.iter} # 'trg_feat': trg_feat
    # [1, 81, 64, 64] [1, 2, 64, 64] [1, 64, 64, 64]
    if settings.env.use_gt_mask == False:
        model_kwsettings['mask_y512'] = seg_map_all # [b, 384, 64, 64]
    if settings.env.use_line_mask == True:
        model_kwsettings['line_msk'] = textline_map # 
    image_size_h, image_size_w = feature_size, feature_size
    # tv_save_image(source,"vis_hp/debug_vis/source.png")
    # tv_save_image(doc_mask,"vis_hp/debug_vis/mask512_8877.png")
    logger.info(f"\nStarting sampling")

    sample, _ = diffusion.ddim_sample_loop(
        model,
        (1, 2, image_size_h, image_size_w), # 1,2,64,64
        noise=None,
        clip_denoised=settings.env.clip_denoised, # false
        model_kwargs=model_kwsettings,
        eta=0.0,
        progress=True,
        denoised_fn=None,
        sampling_kwargs={'src_img': source}, # 'trg_img': target
        logger=logger,
        n_batch=settings.env.n_batch,
        time_variant = settings.env.time_variant,
        pyramid=pyramid
    )

    sample = th.clamp(sample, min=-1, max=1)
    return sample



def run_evaluation_docunet(
    settings, logger, val_loader, diffusion: GaussianDiffusion, model, 
    pretrained_dewarp_model,pretrained_line_seg_model=None,pretrained_seg_model=None
):
    os.makedirs(f'vis_hp/{settings.env.eval_dataset_name}/{settings.name}', exist_ok=True)
    # batch_preprocessing = DocBatchPreprocessing(
    #             settings, apply_mask=False, apply_mask_zero_borders=False, sparse_ground_truth=False
    #         )
    batch_preprocessing = None
    pbar = tqdm(enumerate(val_loader), total=len(val_loader))
    pyramid = VGGPyramid(train=False).to(dist_util.dev())
    SIZE = None
    trian_t = []
    for i, data in pbar:
        radius = 4
        raw_corr = None
        image_size = 64 
        data_path = data['path']
        # ref test
        # source_288 = F.interpolate(data['source_image']/255., size=(288), mode='bilinear', align_corners=True).to(dist_util.dev())
        source_288 = F.interpolate(data['source_image'], size=(288), mode='bilinear', align_corners=True).to(dist_util.dev())
        # tv_save_image(data['source_image']/255., "vis_hp/msk5/in{}".format(data['path'][0].split('/')[-1]))
        
        
        if settings.env.time_variant == True:
            init_feat = torch.zeros((data['source_image'].shape[0], 256, image_size, image_size), dtype=torch.float32).to(dist_util.dev()) 
        else:
            init_feat = None
        
        
        with torch.inference_mode():
            ref_bm, mask_x = pretrained_dewarp_model(source_288) # [1,2,288,288] 0~288  0~1
            # base = coords_grid_tensor((288,288)).to(ref_bm.device) # [1, 2, 288, 288]
            # ref_flow = ref_bm - base 
            ref_flow = ref_bm/287.0 # [-1, 1]  # [1,2,288,288]
        if settings.env.use_init_flow: 
            init_flow = F.interpolate(ref_flow, size=(image_size), mode='bilinear', align_corners=True) # [24, 2, 64, 64]                  
        else:
            init_flow = torch.zeros((data['source_image'].shape[0], 2, image_size, image_size), dtype=torch.float32).to(dist_util.dev()) 
        # mask_x = F.interpolate(mask_x, size=(512), mode='bilinear', align_corners=True) # 0-1
        # data['source_image'] = mask_x*data['source_image'].to(dist_util.dev()) # 0-255
        # mask_x_vis = mask_x*data['source_image'].to(dist_util.dev()) # 不存在最优mask阈值策略
        # tv_save_image(mask_x_vis, "vis_hp/msk_wore/{}".format(data['path'][0].split('/')[-1])) # 0~1 (288,288)
        (
            data,
            H_ori, # 512
            W_ori, # 512
            source, # [1, 3, 512, 512] 0-1
            target, # None
            batch_ori, # None
            batch_ori_inter, # None
            source_256,# [1, 3, 256, 256] 0-1
            target_256, # None
            source_vis, # [1, 3, H, W] cpu仅用于可视化
            target_vis, # None
            mask, # [1, 512, 512] 全白
            source_0
        ) = prepare_data(settings, batch_preprocessing, SIZE, data)
        
        with torch.no_grad():
            if settings.env.use_gt_mask == False:
                # ref_bm, mask_x = self.pretrained_dewarp_model(source_288) # [1,2,288,288] bm 0~288 mskx0-256
                mskx, d0, hx6, hx5d, hx4d, hx3d, hx2d, hx1d = pretrained_seg_model(source_288)
                hx6 = F.interpolate(hx6, size=image_size, mode='bilinear', align_corners=False)
                hx5d = F.interpolate(hx5d, size=image_size, mode='bilinear', align_corners=False)
                hx4d = F.interpolate(hx4d, size=image_size, mode='bilinear', align_corners=False)
                hx3d = F.interpolate(hx3d, size=image_size, mode='bilinear', align_corners=False)
                hx2d = F.interpolate(hx2d, size=image_size, mode='bilinear', align_corners=False)
                hx1d = F.interpolate(hx1d, size=image_size, mode='bilinear', align_corners=False)

                seg_map_all = torch.cat((hx6, hx5d, hx4d, hx3d, hx2d, hx1d), dim=1) # [b, 384, 64, 64]
                # tv_save_image(mskx,"vis_hp/debug_vis/mskx.png")
                if settings.env.use_line_mask:
                    textline_map, textline_mask = pretrained_line_seg_model(mskx) # [3, 64, 256, 256]
                    textline_map = F.interpolate(textline_map, size=image_size, mode='bilinear', align_corners=False) #  [3, 64, 64, 64]
            else:
                seg_map_all = None
                textline_map = None
        
        
         
        
        if settings.env.train_VGG:
            c20 = None
            feature_size = image_size
        else:
            feature_size = image_size
            if settings.env.train_mode == 'stage_1_dit_cat' or settings.env.train_mode =='stage_1_dit_cross':
                with th.no_grad():
                    c20  = extract_raw_features_single2(pyramid, source, source_256, feature_size) # [24, 1, 64, 64, 64, 64]
                # 平均互相关，VGG最浅层特征的下采样（512*512->64*64）
            else:
                with th.no_grad():
                    c20  = extract_raw_features_single(pyramid, source, source_256, feature_size) # [24, 1, 64, 64, 64, 64]
                # 平均互相关，VGG最浅层特征的下采样（512*512->64*64）
        
        source_64 = None # F.interpolate(source, size=(feature_size), mode='bilinear', align_corners=True)
        logger.info(f"Starting sampling with VGG Features")
        # init_flow = correlation_to_flow_w_argmax(
        #         raw_corr.view(1, 1, feature_size, feature_size, feature_size, feature_size),
        #         output_shape=(feature_size, feature_size),
        #     )  # B, 2, 64, 64 初始偏移场
        
        import time
        begin_train = time.time()
        sample = run_sample_lr_dewarping(
            settings,
            logger,
            diffusion,
            model,
            radius, # 4
            source, # [B, 3, 512, 512] 0~1
            feature_size, # 64
            raw_corr, # None
            init_flow, # [B, 2, 64, 64]   -1~1
            c20, # # [B, 64, 64, 64]   
            source_64, # None
            pyramid,
            mask_x, #mask_x,  # F.interpolate(mskx, size=(512), mode='bilinear', align_corners=True)[:,:1,:,:] , # mask_x 
            seg_map_all,
            textline_map,
            init_feat
        ) # sample: [1, 2, 64, 64] 偏移量 [-1,1]范围 五步DDIM的结果
        
        
        
        trian_t.append(time.time()-begin_train)  # 从这里宣布结束训练当前epoch

        
        
        
        

        
        # if settings.env.use_sr_net == True: # false
        #     logger.info('Running super resolution')
        #     sample_sr = None
        #     for j in range(1):
        #         batch_ori, sample_sr, init_flow_sr = run_sample_sr(
        #             settings, logger, diffusion_sr, model_sr, pyramid, data, sample, sample_sr
        #         )

        #         sample_ = F.interpolate(sample_sr, size=(H_ori, W_ori), mode='bilinear', align_corners=True)
        #         sample_[:, 0, :, :] = sample_[:, 0, :, :] * W_ori
        #         sample_[:, 1, :, :] = sample_[:, 1, :, :] * H_ori

        #         sample_ = sample_.permute(0, 2, 3, 1)[mask]
        #         batch_ori_ = batch_ori.permute(0, 2, 3, 1)[mask]
        #         epe = th.sum((sample_ - batch_ori_.to(sample_.device)) ** 2, dim=1).sqrt()
        #         logger.info(f'sr iter: {i}, epe: {epe.mean()}')

        #     sample = F.interpolate(sample_sr, size=(H_ori, W_ori), mode='bilinear', align_corners=True)
        #     sample[:, 0, :, :] = sample[:, 0, :, :] * W_ori
        #     sample[:, 1, :, :] = sample[:, 1, :, :] * H_ori
        #     init_flow = F.interpolate(init_flow_sr, size=(H_ori, W_ori), mode='bilinear', align_corners=True)
        #     # init_flow[:, 0, :, :] = init_flow[:, 0, :, :] * W_ori
        #     # init_flow[:, 1, :, :] = init_flow[:, 1, :, :] * H_ori
        # sample = th.mean(sample[0], dim=0, keepdim=True)
        if settings.env.use_sr_net == False:
            sample = F.interpolate(sample, size=(H_ori, W_ori), mode='bilinear', align_corners=True) # [-1,+1] 偏移场
            # sample[:, 0, :, :] = sample[:, 0, :, :] * W_ori
            # sample[:, 1, :, :] = sample[:, 1, :, :] * H_ori 
            base = F.interpolate(coords_grid_tensor((512,512))/511., size=(H_ori, W_ori), mode='bilinear', align_corners=True)
            # sample = ( ((sample + base.to(sample.device)) )*2 - 1 )
            sample = ( ((sample + base.to(sample.device))*1 )*2 - 1 )*0.987 #  (2 * (bm / 286.8) - 1) * 0.99
            ref_flow = None
            if ref_flow is not None: 
                ref_flow = F.interpolate(ref_flow, size=(H_ori, W_ori), mode='bilinear', align_corners=True) # [-1,+1] 偏移场
                # ref_flow[:, 0, :, :] = ref_flow[:, 0, :, :] * W_ori 
                # ref_flow[:, 1, :, :] = ref_flow[:, 1, :, :] * H_ori 
                ref_flow  = (ref_flow + base.to(ref_flow.device))*2 -1
            # init_flow = F.interpolate(init_flow, size=(H_ori, W_ori), mode='bilinear', align_corners=True)
        else:
            raise ValueError("Invalid value")
        
        if settings.env.visualize:
            visualize_dewarping(settings, sample, data, i, source_vis, data_path, ref_flow)

        
        # sample = sample.permute(0, 2, 3, 1)[mask]
        # init_flow[:, 0, :, :] = init_flow[:, 0, :, :] * W_ori
        # init_flow[:, 1, :, :] = init_flow[:, 1, :, :] * H_ori
        # init_flow = init_flow.permute(0, 2, 3, 1)[mask]
    # print("Elapsed time:{:.2f} minutes ".format(trian_t/60))
    print(len(trian_t))
    print("Elapsed time:{:.2f} avg_second ".format(sum(trian_t) / len(trian_t)))


def coords_grid_tensor(perturbed_img_shape):
    im_x, im_y = np.mgrid[0:perturbed_img_shape[0]-1:complex(perturbed_img_shape[0]),
                            0:perturbed_img_shape[1]-1:complex(perturbed_img_shape[1])]
    coords = np.stack((im_y,im_x), axis=2) # 先x后y，行序优先
    coords = th.from_numpy(coords).float().permute(2,0,1).to(dist_util.dev())  # (2, 512, 512)
    return coords.unsqueeze(0) # [2, 512, 512]


def validate(local_rank, args, val_loader, model, criterion):
    for i, sample in enumerate(val_loader):
        input1, label = sample # [2, 3, 288, 288],[2, 2, 288, 288]
        input1 = input1.to(local_rank,non_blocking=True)
        label = label.to(local_rank,non_blocking=True)
        # label = (label/288.0-0.5)*2
        
        with torch.no_grad():
            output = model(input1) # [3b, 2, 288, 288]
            # loss = F.l1_loss(output, label) #  合成图像强监督
            # test point
            # bm_test=(output/288.0-0.5)*2
            bm_test = (output/992.0-0.5)*2
            label = (label/992.0-0.5)*2
            # bm_test = output
            bm_test = F.interpolate(bm_test, size=(1000,1000), mode='bilinear', align_corners=True)
            label = F.interpolate(label, size=(1000,1000), mode='bilinear', align_corners=True)
            input1 =  F.interpolate(input1, size=(1000,1000), mode='bilinear', align_corners=True)
            regis_image1 = F.grid_sample(input=input1, grid=bm_test.permute(0,2,3,1), align_corners=True)
            regis_image2 = F.grid_sample(input=input1, grid=label.permute(0,2,3,1), align_corners=True)
            
            # regis_image2 = F.grid_sample(input=a_sample[None], grid=bm_test[None].permute(0,2,3,1), align_corners=True)
            tv_save_image(input1[0], "backup/test/ori.png")
            tv_save_image(regis_image1[0], "backup/test/aaa.png")
            tv_save_image(regis_image2[0], "backup/test/gt.png")
            
            # warped_src = warp(source_vis.to(sample.device).float(), sample) # [1, 3, 1629, 981]
            # warped_src = warped_src[0].permute(1, 2, 0).detach().cpu().numpy() # (1873, 1353, 3)
            # warped_src = Image.fromarray((warped_src).astype(np.uint8))
            # warped_src.save(f"vis_hp/{settings.env.eval_dataset_name}/{settings.name}/dewarped_pred/warped_{data_path[0].split('/')[-1]}")



    return None