Update app.py

app.py

@@ -68,12 +68,11 @@ def pred_maps(images):
     with torch.no_grad():
         scaled_preds_tensor = model(images_proc.to(device))[-1]
     preds = []
-    print('scaled_preds_tensor.sum():', scaled_preds_tensor.sum())
-    print('type(scaled_preds_tensor):', type(scaled_preds_tensor))
     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 = cv2.cvtColor((pred_tensor*255).astype(np.uint8)
         cv2.imwrite(save_path, pred_tensor)
 
     zip_file_path = os.path.join(save_dir, "{}.zip".format(save_dir))