ARTeLab/DTM_Estimation_SRandD

1import gradio as gr2import os3from PIL import Image4import torchvision5from torchvision import transforms6import torch7import matplotlib.pyplot as plt8import numpy as np9from models.modelNetA import Generator as GA10from models.modelNetB import Generator as GB11from models.modelNetC import Generator as GC12 13scale_size = 12814scale_sizes = [128, 256, 512]15# load model16modeltype2path = {17    'ModelA': 'DTM_exp_train10%_model_a/g-best.pth',18    'ModelB': 'DTM_exp_train10%_model_b/g-best.pth',19    'ModelC': 'DTM_exp_train10%_model_c/g-best.pth',20}21DEVICE='cpu'22MODELS_TYPE = list(modeltype2path.keys())23generators = [GA(), GB(), GC()]24 25for i in range(len(generators)):26    generators[i] = torch.nn.DataParallel(generators[i])27    state_dict = torch.load(modeltype2path[MODELS_TYPE[i]], map_location=torch.device('cpu'))28    generators[i].load_state_dict(state_dict)29    generators[i] = generators[i].module.to(DEVICE)30    generators[i].eval()31 32preprocess = transforms.Compose([33    transforms.Grayscale(),34    transforms.ToTensor()35])36 37def predict(input_image, model_name, input_scale_factor):38    pil_image = Image.fromarray(input_image.astype('uint8'), 'RGB')39    pil_image = transforms.Resize((input_scale_factor, input_scale_factor))(pil_image)40    # transform image to torch and do preprocessing41    torch_img = preprocess(pil_image).to(DEVICE).unsqueeze(0).to(DEVICE)42    torch_img = (torch_img - torch.min(torch_img)) / (torch.max(torch_img) - torch.min(torch_img))43    # model predict44    with torch.no_grad():45        output = generators[MODELS_TYPE.index(model_name)](torch_img)46    sr, sr_dem_selected = output[0], output[1]47    # transform torch to image48    sr = sr.squeeze(0).cpu()49    torchvision.utils.save_image(sr, 'sr_pred.png')50    sr = np.array(Image.open('sr_pred.png'))51 52    sr_dem_selected = sr_dem_selected.squeeze().cpu().detach().numpy()53    fig, ax = plt.subplots()54    im = ax.imshow(sr_dem_selected, cmap='jet', vmin=0, vmax=np.max(sr_dem_selected))55    plt.colorbar(im, ax=ax)56    fig.canvas.draw()57    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)58    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))59    # return correct image and info60    info = f"{model_name} with {sum(p.numel() for p in generators[MODELS_TYPE.index(model_name)].parameters())} parameters"61    return info, sr, data62 63iface = gr.Interface(64    fn=predict, 65    inputs=[66        gr.Image(),67        gr.inputs.Radio(MODELS_TYPE),68        gr.inputs.Radio(scale_sizes)69    ], 70    outputs=[71        gr.Text(label='Model info'),72        gr.Image(label='Super Resolution'),73        gr.Image(label='DTM')74    ],75    examples=[76        [f"demo_imgs/{name}", MODELS_TYPE[0], 128] for name in os.listdir('demo_imgs')77    ],78    title="Super Resolution and DTM Estimation",79    description=f"This demo predict Super Resolution and (Super Resolution) DTM from a Grayscale image (if RGB we convert it)."80)81iface.launch()