Update app.py

app.py

@@ -72,10 +72,10 @@ 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,wl,hl):
+def test(gpu_id, net, img_list, group_size, img_size):
     print('test')
     #device=device
-    
+    hl,wl=[_.shape[0] for _ in img_list],[_.shape[1] for _ in img_list]
     img_transform = transforms.Compose([transforms.Resize((img_size, img_size)), transforms.ToTensor(),
                     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
     img_transform_gray = transforms.Compose([transforms.Resize((img_size, img_size)), transforms.ToTensor(),
@@ -113,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,w_list,h_list)
+  result_list=test(device,net,img_list,5,224)
   #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