Update app.py

app.py

@@ -29,6 +29,14 @@ class ImagePreprocessor():
         return image
 
 
+def save_tensor_img(path, tenor_im):
+    im = tenor_im.cpu().clone()
+    im = im.squeeze(0)
+    tensor2pil = transforms.ToPILImage()
+    im = tensor2pil(im)
+    im.save(path)
+
+
 model = GCoNet(bb_pretrained=False).to(device)
 state_dict = './ultimate_duts_cocoseg (The best one).pth'
 if os.path.exists(state_dict):
@@ -71,9 +79,8 @@ def pred_maps(images):
     for image_shape, pred_tensor, save_path in zip(image_shapes, scaled_preds_tensor, save_paths):
         if device == 'cuda':
             pred_tensor = pred_tensor.cpu()
-        pred_tensor = torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze().numpy()
-        pred_tensor = (pred_tensor*255).astype(np.uint8)
-        cv2.imwrite(save_path, pred_tensor)
+        pred_tensor = torch.nn.functional.interpolate(pred_tensor.unsqueeze(0), size=image_shape, mode='bilinear', align_corners=True).squeeze()
+        save_tensor_img(save_path, pred_tensor)
 
     zip_file_path = os.path.join(save_dir, "{}.zip".format(save_dir))
     with zipfile.ZipFile(zip_file_path, 'w') as zipf:
@@ -82,9 +89,6 @@ def pred_maps(images):
     return save_paths, zip_file_path
 
 
-# N = 4
-# examples = [[_] for _ in glob('example_images/butterfly/*')][:N]
-
 tab_batch = gr.Interface(
     fn=pred_maps,
     inputs=gr.File(label="Upload multiple images in a group", type="filepath", file_count="multiple"),