Update app.py

app.py

@@ -13,6 +13,7 @@ os.system("pip3 install git+https://github.com/lucasb-eyer/pydensecrf.git")
 import pydensecrf.densecrf as dcrf
 from PIL import Image
 import torch
+import torch.nn.functional as F
 from torchvision import transforms
 from model_video import build_model
 import numpy as np
@@ -71,7 +72,7 @@ for k in weight.keys():
 net.load_state_dict(new_dict)
 net.eval()
 net = net.to(device)
-def test(gpu_id, net, img_list, group_size, img_size):
+def test(gpu_id, net, img_list, group_size, img_size,wl,hl):
     print('test')
     #device=device
     
@@ -85,8 +86,10 @@ def test(gpu_id, net, img_list, group_size, img_size):
         for i in range(5):
           group_img[i]=img_transform(Image.fromarray(img_list[i]))
         _,pred_mask=net(group_img*1)
-        pred_mask=(pred_mask.detach().squeeze()*255).numpy().astype(np.uint8)
-        pred_mask=[crf_refine(((group_img[i]-group_img[i].min())/(group_img[i].max()-group_img[i].min())*255).permute(1,2,0).contiguous().numpy().astype(np.uint8),pred_mask[i]) for i in range(5)]
+        pred_mask=(pred_mask.detach().squeeze()*255)#.numpy().astype(np.uint8)
+        pred_mask=[F.interpolate(pred_mask[i].reshape(1,1,pred_mask[i].shape[-2],pred_mask[i].shape[-1]),size=(wl[i],hl[i]),mode='bilinear').numpy().astype(np.uint8) for i in range(5)]
+        #pred_mask=[crf_refine(((group_img[i]-group_img[i].min())/(group_img[i].max()-group_img[i].min())*255).permute(1,2,0).contiguous().numpy().astype(np.uint8),pred_mask[i]) for i in range(5)]
+        pred_mask=[crf_refine(img_list[i],pred_mask[i]) for i in range(5)]
         print(pred_mask[0].shape)
         result = [Image.fromarray((torch.from_numpy(pred_mask[i]).unsqueeze(2).repeat(1,1,3)).numpy()) for i in range(5)]
         #w, h = 224,224#Image.open(image_list[i][j]).size
@@ -110,7 +113,7 @@ def sepia(img1,img2,img3,img4,img5):
   h_list,w_list=[_.shape[0] for _ in img_list],[_.shape[1] for _ in img_list]
   #print(type(img1))
   #print(img1.shape)
-  result_list=test(device,net,img_list,5,224)
+  result_list=test(device,net,img_list,5,224,w_list,h_list)
   #result_list=[result_list[i].resize((w_list[i], h_list[i]), Image.BILINEAR) for i in range(5)]
   img1,img2,img3,img4,img5=result_list#test('cpu',net,img_list,5,224)
   return img1,img2,img3,img4,img5