added img size parameters

qim3d/viz/img.py

@@ -192,13 +192,13 @@ def grid_pred(in_targ_preds, num_images=7, cmap_im="gray", cmap_segm="viridis",
 
     fig.show()
 
-def slice_viz(input, position = 'mid', cmap_im="viridis",axis = False):
+def slice_viz(input, position = 'mid', cmap="viridis", axis=False, img_height=2, img_width=2):
     """ Displays one or several slices from a 3d array.
 
     Args:
         input (str, numpy.ndarray): Path to the file or 3-dimensional array.
         position (str, int, list, array, optional): One or several slicing levels.
-        cmap_im (str, optional): Specifies the color map for the image.
+        cmap (str, optional): Specifies the color map for the image.
         axis (bool, optional): Specifies whether the axes should be included.
 
     Raises:
@@ -233,12 +233,12 @@ def slice_viz(input, position = 'mid', cmap_im="viridis",axis = False):
     if isinstance(position,str):
         if position.lower() in ['mid','middle']:
             height = [int(vol.shape[-1]/2)]
-        elif position.lower() in ['top','upper']:
+        elif position.lower() in ['top','upper', 'start']:
             height = [0]
-        elif position.lower() in ['bot','bottom']:
+        elif position.lower() in ['bot','bottom', 'end']:
             height = [vol.shape[-1]-1]
         else:
-            raise ValueError('Position not recognized. Choose an integer, list, array or "top","mid","bot".')
+            raise ValueError('Position not recognized. Choose an integer, list, array or "start","mid","end".')
     
     # Position is an integer
     elif isinstance(position,int):
@@ -250,12 +250,12 @@ def slice_viz(input, position = 'mid', cmap_im="viridis",axis = False):
 
     num_images = len(height)
     
-    fig = plt.figure(figsize=(2 * num_images, 5), constrained_layout = True)
+    fig = plt.figure(figsize=(img_width * num_images, img_height), constrained_layout = True)
     axs = fig.subplots(nrows = 1, ncols = num_images)
     
     for col, ax in enumerate(np.atleast_1d(axs)):
-        ax.imshow(vol[:,:,height[col]],cmap = cmap_im)
-        ax.title.set_text(f'Slice {height[col]}')
+        ax.imshow(vol[:,:,height[col]],cmap = cmap)
+        ax.set_title(f'Slice {height[col]}', fontsize=8)
         if not axis:
             ax.axis('off')