diff --git a/pysegcnn/core/predict.py b/pysegcnn/core/predict.py
index 6403a7aed9b20dedd9bd914d075377e97fad91f9..ff7455e931e02e32a840b712709445ae0706df17 100644
--- a/pysegcnn/core/predict.py
+++ b/pysegcnn/core/predict.py
@@ -245,9 +245,6 @@ def predict_scenes(ds, model, scene_id=None, cm=False, plot=False, **kwargs):
         # the name of the selected scene
         scene_ids = [scene_id]
 
-    # spatial size of scene
-    scene_size = (ds.dataset.height, ds.dataset.width)
-
     # iterate over the scenes
     LOGGER.info('Predicting scenes of the {} dataset ...'.format(ds.name))
     output = {}
@@ -284,9 +281,9 @@ def predict_scenes(ds, model, scene_id=None, cm=False, plot=False, **kwargs):
                     conf_mat[ytrue.long(), ypred.long()] += 1
 
         # reconstruct the entire scene
-        inputs = reconstruct_scene(inp, scene_size, nbands=inp.shape[1])
-        labels = reconstruct_scene(lab, scene_size, nbands=1)
-        prdtcn = reconstruct_scene(prd, scene_size, nbands=1)
+        inputs = reconstruct_scene(inp)
+        labels = reconstruct_scene(lab)
+        prdtcn = reconstruct_scene(prd)
 
         # print progress
         LOGGER.info('Scene {:d}/{:d}, Id: {}, Accuracy: {:.2f}'.format(
diff --git a/pysegcnn/core/utils.py b/pysegcnn/core/utils.py
index 59d08f9c73d8d6e90137b66546c0dd9e9dfb29eb..1729a4aec120bdf7958f923b9fac30577ce51d34 100644
--- a/pysegcnn/core/utils.py
+++ b/pysegcnn/core/utils.py
@@ -398,19 +398,14 @@ def tile_topleft_corner(img_size, tile_size):
     return indices
 
 
-def reconstruct_scene(tiles, img_size, tile_size=None, nbands=1):
+def reconstruct_scene(tiles):
     """Reconstruct a tiled image.
 
     Parameters
     ----------
-    tiles : array_like
-        The tiled image, shape=(tiles, bands, tile_size, tile_size).
-    img_size : `tuple`
-        The size of the reconstructed image (height, width).
-    tile_size : `int` or `None`, optional
-        The size of the tile. The default is None.
-    nbands : `int`, optional
-        The number of bands of the reconstructed image. The default is 1.
+    tiles : `torch.Tensor` or `numpy.ndarray`
+        The tiled image, shape=(tiles, bands, tile_size, tile_size) or
+        shape=(tiles, tile_size, tile_size).
 
     Returns
     -------
@@ -422,11 +417,15 @@ def reconstruct_scene(tiles, img_size, tile_size=None, nbands=1):
     tiles = np.asarray(tiles)
 
     # check the size
-    if tile_size is None:
-        if tiles.ndim > 3:
-            tile_size = tiles.shape[2]
-        else:
-            tile_size = tiles.shape[1]
+    if tiles.ndim > 3:
+        nbands = tiles.shape[1]
+        tile_size = tiles.shape[2]
+    else:
+        nbands = 1
+        tile_size = tiles.shape[1]
+
+    # calculate image 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)