app use models

.gitignore

@@ -2,4 +2,6 @@ venv
 *.pyc
 __pycache__
 sr.png
+sr2.png
+sr_pred.png
 test.png

app.py

@@ -1,21 +1,58 @@
 import gradio as gr
 from PIL import Image
 import torchvision
+from torchvision import transforms
 import torch
+import matplotlib.pyplot as plt
+import numpy as np
+from models.modelNetA import Generator as GA
+from models.modelNetB import Generator as GB
+from models.modelNetC import Generator as GC
 
 # load model
-MODELS_TYPE = ["ModelA", "ModelB", "ModelC"]
+modeltype2path = {
+    'ModelA': 'DTM_exp_train10%_model_a/g-best.pth',
+    'ModelB': 'DTM_exp_train10%_model_b/g-best.pth',
+    'ModelC': 'DTM_exp_train10%_model_c/g-best.pth',
+}
+DEVICE='cpu'
+MODELS_TYPE = list(modeltype2path.keys())
+generators = [GA(), GB(), GC()]
+
+for i in range(len(generators)):
+    generators[i] = torch.nn.DataParallel(generators[i])
+    state_dict = torch.load(modeltype2path[MODELS_TYPE[i]], map_location=torch.device('cpu'))
+    generators[i].load_state_dict(state_dict)
+    generators[i] = generators[i].module.to(DEVICE)
+    generators[i].eval()
+
+preprocess = transforms.Compose([
+    transforms.Grayscale(),
+    transforms.ToTensor()
+])
 
 def predict(input_image, model_name):
     pil_image = Image.fromarray(input_image.astype('uint8'), 'RGB')
     # transform image to torch and do preprocessing
-    torch_image = torchvision.transforms.ToTensor()(pil_image)
+    torch_img = preprocess(pil_image).to(DEVICE).unsqueeze(0).to(DEVICE)
     # model predict
-    prediction = torch.rand(torch_image.shape)
+    with torch.no_grad():
+        output = generators[MODELS_TYPE.index(model_name)](torch_img)
+    sr, sr_dem_selected = output[0], output[1]
     # transform torch to image
-    predicted_pil_image = torchvision.transforms.ToPILImage()(prediction)
+    sr = sr.squeeze(0).cpu()
+    torchvision.utils.save_image(sr, 'sr_pred.png')
+    sr = np.array(Image.open('sr_pred.png'))
+
+    sr_dem_selected = sr_dem_selected.squeeze().cpu().detach().numpy()
+    fig, ax = plt.subplots()
+    im = ax.imshow(sr_dem_selected, cmap='jet', vmin=0, vmax=np.max(sr_dem_selected))
+    plt.colorbar(im, ax=ax)
+    fig.canvas.draw()
+    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
     # return correct image
-    return predicted_pil_image
+    return sr, data
 
 iface = gr.Interface(
     fn=predict, 
@@ -23,7 +60,10 @@ iface = gr.Interface(
         gr.Image(shape=(512,512)),
          gr.inputs.Radio(MODELS_TYPE)
     ], 
-    outputs=gr.Image(shape=(512,512)),
+    outputs=[
+        gr.Image(),
+        gr.Image()
+    ],
     examples=[
         ["demo_imgs/fake.jpg", MODELS_TYPE[0]] # use real image
     ],

test.py

@@ -35,7 +35,7 @@ generator.eval()
 
 preprocess = transforms.Compose([
     transforms.Grayscale(),
-    transforms.Resize((512, 512)),
+    # transforms.Resize((512, 512)),
     transforms.ToTensor()
 ])
 input_img = Image.open('demo_imgs/fake.jpg')
@@ -47,6 +47,8 @@ sr = sr.squeeze(0).cpu()
 
 print(sr.shape)
 torchvision.utils.save_image(sr, 'sr.png')
+# sr = Image.fromarray(sr.squeeze(0).detach().numpy() * 255, 'L')
+# sr.save('sr2.png')
 
 sr_dem_selected = sr_dem_selected.squeeze().cpu().detach().numpy()
 print(sr_dem_selected.shape)