From b206efdaa729accb1f9ed219b904ee96cea4ac62 Mon Sep 17 00:00:00 2001
From: "Daniel.Frisinghelli" <daniel.frisinghelli@eurac.edu>
Date: Fri, 5 Feb 2021 15:33:59 +0100
Subject: [PATCH] Moved check for gpu to predict method of NetworkInference.
---
pysegcnn/core/trainer.py | 9 +++++++--
pysegcnn/core/utils.py | 22 ++++++++--------------
2 files changed, 15 insertions(+), 16 deletions(-)
diff --git a/pysegcnn/core/trainer.py b/pysegcnn/core/trainer.py
index c3066ad..f3089ff 100644
--- a/pysegcnn/core/trainer.py
+++ b/pysegcnn/core/trainer.py
@@ -2513,6 +2513,11 @@ class NetworkInference(BaseConfig):
progress = 'Sample: {:d}/{:d}'.format(batch + 1,
len(self.dataloader))
+ # check if tensor is on gpu and convert to numpy array
+ inputs = inputs.cpu().numpy()
+ labels = labels.cpu().numpy()
+ prdctn = prdctn.cpu().numpy()
+
# check whether to reconstruct the scene
if self.dataloader.batch_size > 1:
@@ -2643,9 +2648,9 @@ class NetworkInference(BaseConfig):
if self.cm:
# merge predictions for all samples
- y_true = np.asarray([v['y_true'].numpy().flatten() for _, v
+ y_true = np.asarray([v['y_true'].flatten() for _, v
in output.items()]).flatten()
- y_pred = np.asarray([v['y_pred'].numpy().flatten() for _, v
+ y_pred = np.asarray([v['y_pred'].flatten() for _, v
in output.items()]).flatten()
# calculate confusion matrix
diff --git a/pysegcnn/core/utils.py b/pysegcnn/core/utils.py
index 24d905b..609deda 100644
--- a/pysegcnn/core/utils.py
+++ b/pysegcnn/core/utils.py
@@ -745,7 +745,7 @@ def reconstruct_scene(tiles):
Parameters
----------
- tiles : :py:class:`torch.Tensor` or :py:class:`numpy.ndarray`
+ tiles : :py:class:`numpy.ndarray`
The tiled image, shape: `(tiles, bands, tile_size, tile_size)` or
`(tiles, tile_size, tile_size)`.
@@ -755,32 +755,26 @@ def reconstruct_scene(tiles):
The reconstructed image, shape: `(bands, height, width)`.
"""
- # check if tensor is on gpu and convert to numpy array
- if isinstance(tiles, torch.Tensor):
- tiles_cpu = np.asarray(tiles.cpu())
- else:
- tiles_cpu = np.asarray(tiles)
-
# check the dimensions of the input array
- if tiles_cpu.ndim > 3:
- nbands = tiles_cpu.shape[1]
- tile_size = tiles_cpu.shape[2]
+ if tiles.ndim > 3:
+ nbands = tiles.shape[1]
+ tile_size = tiles.shape[2]
else:
nbands = 1
- tile_size = tiles_cpu.shape[1]
+ tile_size = tiles.shape[1]
# calculate image size
- img_size = 2 * (int(np.sqrt(tiles_cpu.shape[0]) * tile_size),)
+ img_size = 2 * (int(np.sqrt(tiles.shape[0]) * tile_size),)
# calculate the topleft corners of the tiles
topleft = tile_topleft_corner(img_size, tile_size)
# iterate over the tiles
scene = np.zeros(shape=(nbands,) + img_size)
- for t in range(tiles_cpu.shape[0]):
+ for t in range(tiles.shape[0]):
scene[...,
topleft[t][0]: topleft[t][0] + tile_size,
- topleft[t][1]: topleft[t][1] + tile_size] = tiles_cpu[t, ...]
+ topleft[t][1]: topleft[t][1] + tile_size] = tiles[t, ...]
return scene.squeeze()
--
GitLab