Update app.py

[fix] dont cat when single

app.py

@@ -124,7 +124,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,stack_image=True):
     print('test')
     #device=device
     hl,wl=[_.shape[0] for _ in img_list],[_.shape[1] for _ in img_list]
@@ -143,6 +143,8 @@ def test(gpu_id, net, img_list, group_size, img_size):
         img_resize=[((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) 
                     for i in range(5)]
         pred_mask=[(pred_mask[i].numpy().astype(np.uint8)) for i in range(5)]#[(img_resize[i],pred_mask[i].numpy().astype(np.uint8)) for i in range(5)]
+        if not stack_image:
+            return pred_mask[0]
         #for i in range(5):
         #    print(img_list[i].shape,pred_mask[i].shape)
         #pred_mask=[crf_refine(img_list[i],pred_mask[i]) for i in range(5)]
@@ -177,12 +179,13 @@ def sepia(img1,img2,img3,img4,img5,stack_image=True):
   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,stack_image)
+  if not stack_image:
+     return result_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)
   white=(torch.ones(img1.shape[0],2,3)*255).numpy().astype(np.uint8)
-  if not stack_image:
-     return img1
+  
   return np.concatenate([img1,white,img2,white,img3,white,img4,white,img5],axis=1)
 
 #gr.Image(shape=(224, 2))