Add application file

.gitignore

@@ -3,7 +3,7 @@ vis_hp
 assets
 images
 backup
-run_gradio.py
+# run_gradio.py
 run_foward.py
 
 

run_gradio.py

@@ -0,0 +1,553 @@
+import argparse
+import random
+from datetime import date
+from shutil import copyfile
+import cv2 as cv
+import numpy as np
+import torch
+import torch.backends.cudnn
+import admin.settings as ws_settings
+import os
+import torch
+import torch.distributed as dist
+import torchvision.transforms as transforms
+from torch.utils.data import DataLoader
+import datasets
+from utils_data.image_transforms import ArrayToTensor
+from train_settings.dvd.improved_diffusion import dist_util, logger
+from train_settings.dvd.improved_diffusion.script_util import args_to_dict, create_model_and_diffusion,model_and_diffusion_defaults
+from train_settings.models.geotr.geotr_core import GeoTr_Seg_Inf, reload_segmodel, reload_model, Seg   
+from train_settings.models.geotr.unet_model import UNet
+from PIL import Image
+from tqdm import tqdm
+import torch.nn.functional as F
+import torch as th
+from train_settings.dvd.improved_diffusion.gaussian_diffusion import GaussianDiffusion
+from train_settings.dvd.feature_backbones.VGG_features import VGGPyramid
+from train_settings.dvd.eval_utils import extract_raw_features_single,extract_raw_features_single2
+from datasets.utils.warping import register_model2
+
+import gradio as gr
+
+
+
+
+reg_model_bilin = register_model2((512,512), 'bilinear')
+
+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 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
+):
+    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
+
+    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 visualize_dewarping(settings, sample, data, i, source_vis, data_path, ref_flow=None):
+    os.makedirs(f'vis_hp/{settings.env.eval_dataset_name}/{settings.name}/dewarped_pred', exist_ok=True) # pred dewarped
+    # warped_src = warp(source_vis.to(sample.device).float(), sample) # [1, 3, 1629, 981]
+    warped_src = reg_model_bilin([source_vis.to(sample.device).float(), sample])
+    warped_src = warped_src[0].permute(1, 2, 0).detach().cpu().numpy()#*255. # (1873, 1353, 3)
+    warped_src = Image.fromarray((warped_src).astype(np.uint8))
+
+    return warped_src
+
+def visualize_dewarping_single(settings, sample, source_vis):
+    os.makedirs(f'vis_hp/{settings.env.eval_dataset_name}/{settings.name}/dewarped_pred', exist_ok=True) # pred dewarped
+    # warped_src = warp(source_vis.to(sample.device).float(), sample) # [1, 3, 1629, 981]
+    warped_src = reg_model_bilin([source_vis.to(sample.device).float(), sample])
+    warped_src = warped_src[0].permute(1, 2, 0).detach().cpu().numpy()#*255. # (1873, 1353, 3)
+    warped_src = Image.fromarray((warped_src).astype(np.uint8))
+
+    return warped_src
+
+
+
+
+
+
+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
+
+    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
+    
+    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 prepare_data_single(input_image, input_image_ori):
+    source_vis = input_image_ori
+    target_vis = None
+    _, _, H_ori, W_ori = source_vis.shape
+    source = input_image.to(dist_util.dev())  # [1, 3, 914, 1380]  torch.float32
+    source_0 = None
+    target = None 
+    batch_ori = None    
+    batch_ori_inter = None 
+    target = None
+    target_256 = None
+    source_256 = torch.nn.functional.interpolate(input=source.float(), size=(256, 256), mode='area').to(dist_util.dev())
+    target_256 = None
+    mask = torch.ones((1, 512, 512), dtype=torch.bool).to(dist_util.dev()) # None
+
+    return input_image, H_ori, W_ori, source, target, batch_ori, batch_ori_inter, source_256, target_256, source_vis, target_vis, mask, source_0
+
+
+
+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 = None
+    pbar = tqdm(enumerate(val_loader), total=len(val_loader))
+    pyramid = VGGPyramid(train=False).to(dist_util.dev())
+    SIZE = None
+
+    # for each document image
+
+    for i, data in pbar:
+        radius = 4
+        raw_corr = None
+        data_path = data['path']
+        source_288 = F.interpolate(data['source_image'], size=(288), mode='bilinear', align_corners=True).to(dist_util.dev())
+
+        if settings.env.time_variant == True:
+            init_feat = torch.zeros((data['source_image'].shape[0], 256, 64, 64), 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
+            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=(64), mode='bilinear', align_corners=True) # [24, 2, 64, 64]                  
+        else:
+            init_flow = torch.zeros((data['source_image'].shape[0], 2, 64, 64), dtype=torch.float32).to(dist_util.dev()) 
+
+
+        (
+            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=64, mode='bilinear', align_corners=False)
+                hx5d = F.interpolate(hx5d, size=64, mode='bilinear', align_corners=False)
+                hx4d = F.interpolate(hx4d, size=64, mode='bilinear', align_corners=False)
+                hx3d = F.interpolate(hx3d, size=64, mode='bilinear', align_corners=False)
+                hx2d = F.interpolate(hx2d, size=64, mode='bilinear', align_corners=False)
+                hx1d = F.interpolate(hx1d, size=64, 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=64, 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 = 64
+        else:
+            feature_size = 64
+            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")
+
+        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的结果
+        
+
+        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:
+            output = visualize_dewarping(settings, sample, data, i, source_vis, data_path, ref_flow)
+        
+        
+
+def run_single_docunet(input_image_ori):
+    input_image_ori = np.array(input_image_ori, dtype=np.uint8)  # [x, y, 3]
+
+    # resize to 512x512
+    input_image_resized = cv.resize(input_image_ori, (512, 512))  # [512, 512, 3]
+
+    # transpose to [3, 512, 512]
+    input_image_ori = np.transpose(input_image_ori, (2, 0, 1))  # [3, 512, 512]
+    input_image = np.transpose(input_image_resized, (2, 0, 1))  # [3, 512, 512]
+
+    input_image = input_image / 255
+
+    input_image_ori = torch.tensor(input_image_ori).unsqueeze(0) # [1, 3, 512, 512]
+    input_image = torch.tensor(input_image).unsqueeze(0).float() # [1, 3, 512, 512]
+
+    os.makedirs(f'vis_hp/{settings.env.eval_dataset_name}/{settings.name}', exist_ok=True)
+    batch_preprocessing = None
+    pyramid = VGGPyramid(train=False).to(dist_util.dev())
+    SIZE = None
+
+
+    radius = 4
+    raw_corr = None
+    source_288 = F.interpolate(input_image, size=(288), mode='bilinear', align_corners=True).to(dist_util.dev())
+
+    if settings.env.time_variant == True:
+        init_feat = torch.zeros((input_image.shape[0], 256, 64, 64), 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
+        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=(64), mode='bilinear', align_corners=True) # [24, 2, 64, 64]                  
+    else:
+        init_flow = torch.zeros((input_image.shape[0], 2, 64, 64), dtype=torch.float32).to(dist_util.dev()) 
+
+    (
+        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_single(input_image, input_image_ori)
+
+
+
+    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=64, mode='bilinear', align_corners=False)
+            hx5d = F.interpolate(hx5d, size=64, mode='bilinear', align_corners=False)
+            hx4d = F.interpolate(hx4d, size=64, mode='bilinear', align_corners=False)
+            hx3d = F.interpolate(hx3d, size=64, mode='bilinear', align_corners=False)
+            hx2d = F.interpolate(hx2d, size=64, mode='bilinear', align_corners=False)
+            hx1d = F.interpolate(hx1d, size=64, 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=64, 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 = 64
+    else:
+        feature_size = 64
+        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")
+
+    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的结果
+
+    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")
+
+
+    output = visualize_dewarping_single(settings, sample, source_vis)
+
+    return output
+
+
+
+
+
+
+
+
+
+
+parser = argparse.ArgumentParser(description='Run a sampling scripts in train_settings.')
+parser.add_argument('--train_module', type=str, default='dvd', help='Name of module in the "train_settings/" folder.')
+parser.add_argument('--train_name', type=str, default='val_TDiff', help='Name of the train settings file.')
+parser.add_argument('--cudnn_benchmark', type=bool, default=True, help='Set cudnn benchmark on (1) or off (0) (default is on).')
+parser.add_argument('--seed', type=int, default=1992, help='Pseudo-RNG seed')
+parser.add_argument('--name', type=str, default="gradio", help='Name of the experiment')
+parser.add_argument('--corruption', action='store_true') # 默认为false，触发则为true
+
+args = parser.parse_args()
+
+args.seed = random.randint(0, 3000000)
+args.seed = torch.initial_seed() & (2 ** 32 - 1)
+print('Seed is {}'.format(args.seed))
+random.seed(int(args.seed))
+np.random.seed(args.seed)
+    
+cudnn_benchmark=args.cudnn_benchmark
+seed=args.seed
+corruption=args.corruption
+name=args.name
+
+# This is needed to avoid strange crashes related to opencv
+cv.setNumThreads(0)
+
+torch.backends.cudnn.benchmark = cudnn_benchmark
+
+# dd/mm/YY
+today = date.today()
+d1 = today.strftime("%d/%m/%Y")
+print('Sampling:  {}  {}\nDate: {}'.format(args.train_module, args.train_name, d1))
+
+settings = ws_settings.Settings()
+settings.module_name = args.train_module
+settings.script_name = args.train_name
+settings.project_path = 'train_settings/{}/{}'.format(args.train_module, args.train_name) # 'train_settings/DiffMatch/val_DiffMatch'
+settings.seed = seed
+settings.name = name
+
+save_dir = os.path.join(settings.env.workspace_dir, settings.project_path) # 'checkpoints+train_settings/DiffMatch/val_DiffMatch'
+if not os.path.exists(save_dir):
+    os.makedirs(save_dir)
+copyfile(settings.project_path + '.py', os.path.join(save_dir, settings.script_name + '.py'))
+
+
+settings.severity = 0
+settings.corruption_number = 0
+
+
+dist_util.setup_dist()
+logger.configure(dir=f"SAMPLING_{settings.env.eval_dataset}_{settings.name}")
+logger.log(f"Corruption Disabled. Evaluating on Original {settings.env.eval_dataset}")
+logger.log("Loading model and diffusion...")
+
+model, diffusion = create_model_and_diffusion(
+    device=dist_util.dev(),
+    train_mode=settings.env.train_mode, # stage 1
+    tv=settings.env.time_variant,
+    **args_to_dict(settings, model_and_diffusion_defaults().keys()),
+)
+setattr(diffusion, "settings", settings)
+
+pretrained_dewarp_model = GeoTr_Seg_Inf()
+reload_segmodel(pretrained_dewarp_model.msk, settings.env.seg_model_path)
+# reload_model(pretrained_dewarp_model.GeoTr, settings.env.dewarping_model_path)
+pretrained_dewarp_model.to(dist_util.dev())
+pretrained_dewarp_model.eval()
+
+if settings.env.use_line_mask:
+    pretrained_line_seg_model = UNet(n_channels=3, n_classes=1)
+    pretrained_seg_model = Seg()
+    line_model_ckpt = dist_util.load_state_dict(settings.env.line_seg_model_path, map_location='cpu')['model']
+    pretrained_line_seg_model.load_state_dict(line_model_ckpt, strict=True)    
+    pretrained_line_seg_model.to(dist_util.dev())
+    pretrained_line_seg_model.eval()
+
+    seg_model_ckpt = dist_util.load_state_dict(settings.env.new_seg_model_path, map_location='cpu')['model']
+    pretrained_seg_model.load_state_dict(seg_model_ckpt, strict=True)    
+    pretrained_seg_model.to(dist_util.dev())
+    pretrained_seg_model.eval()
+
+model.cpu().load_state_dict(dist_util.load_state_dict(settings.env.model_path, map_location="cpu"), strict=False)
+logger.log(f"Model loaded with {settings.env.model_path}")
+
+model.to(dist_util.dev())
+model.eval()
+
+
+if __name__ == '__main__':
+    demo = gr.Interface(
+        fn=run_single_docunet,
+        inputs=[
+            gr.Image(type="numpy", label="Input Image"),
+        ],
+        outputs=[
+            gr.Image(type="numpy", label="Output Image"),
+        ],
+        title="Document Image Dewarping",
+        description="This is a demo for document image dewarping using a trained model.",
+    )
+
+    demo.launch(share=True, debug=True, server_name="10.7.88.77")