Fixed display issues for Jupyter notebooks in viz module

Author: amogh7joshi

agml/viz/__init__.py

@@ -23,6 +23,7 @@
 from .masks import (
     convert_mask_to_colored_image,
     annotate_semantic_segmentation,
+    show_image_and_mask,
     show_image_with_overlaid_mask,
     show_semantic_segmentation_truth_and_prediction,
 )

agml/viz/boxes.py

@@ -193,7 +193,8 @@ def show_image_and_boxes(image,
                                       **kwargs)
 
     # Display the image.
-    display_image(image)
+    _ = display_image(image, matplotlib_figure = False)
+    return image
 
 
 def show_object_detection_truth_and_prediction(image,
@@ -312,7 +313,7 @@ def show_object_detection_truth_and_prediction(image,
     # Display and return the image.
     image = convert_figure_to_image()
     if not kwargs.get('no_show', False):
-        display_image(image)
+        _ = display_image(image)
     return image
 
 

agml/viz/display.py

@@ -34,17 +34,33 @@ def display_image(image, **kwargs):
             cv2_imshow(image)
             return
 
-        if kwargs.get('read_raw', False):
-            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # convert back to BGR
-        cv2.imshow('image', image)
-        cv2.waitKey(0)
-        cv2.destroyWindow('image')
-
-    if get_viz_backend() == 'matplotlib':
-        plt.figure(figsize = (10, 10))
-        plt.imshow(image)
-        plt.gca().axis('off')
-        plt.gca().set_aspect('equal')
-        plt.show()
+        # If running in a Jupyter notebook, then for some weird reason it automatically
+        # displays images in the background, so don't actually do anything here.
+        notebook = False
+        try:
+            shell = eval("get_ipython().__class__.__name__")
+            if shell == 'ZMQInteractiveShell':
+                notebook = True
+        except NameError:
+            pass
+        if notebook:
+            # If the input content is not a figure, then we can display it.
+            if kwargs.get('matplotlib_figure', True):
+                return
+
+        else:
+            if kwargs.get('read_raw', False):
+                image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR) # convert back to BGR
+            cv2.imshow('image', image)
+            cv2.waitKey(0)
+            cv2.destroyWindow('image')
+            return
+
+    # Default case is matplotlib, since it is the most modular.
+    fig = plt.figure(figsize = (10, 10))
+    plt.imshow(image)
+    plt.gca().axis('off')
+    plt.gca().set_aspect('equal')
+    plt.show()
 
 

agml/viz/labels.py

@@ -57,7 +57,7 @@ def show_images_and_labels(images,
     The matplotlib figure with the plotted info.
     """
     if images is not None and labels is None:
-        if is_array_like(images[0]):
+        if is_array_like(images[0], no_list = True):
             if images[0].ndim >= 3:
                 images, labels = images[0], images[1]
             else:
@@ -82,8 +82,9 @@ def show_images_and_labels(images,
 
     # Check if the labels are converted to one-hot, and re-convert them back.
     if is_array_like(labels):
-        if labels.ndim == 2: # noqa
-            labels = np.argmax(labels, axis = -1)
+        if not isinstance(labels, (list, tuple)):
+            if labels.ndim == 2: # noqa
+                labels = np.argmax(labels, axis = -1)
 
     # If a prime number is passed, e.g. 23, then the `_inference_best_shape`
     # method will return the shape of (23, 1). Likely, the user is expecting
@@ -118,7 +119,7 @@ def show_images_and_labels(images,
     # Display and return the image.
     image = convert_figure_to_image()
     if not kwargs.get('no_show', False):
-        display_image(image)
+        _ = display_image(image)
     return image
 
 

agml/viz/masks.py

@@ -173,7 +173,7 @@ def show_image_with_overlaid_mask(image,
 
     # Display the annotated image.
     if not kwargs.get('no_show', False):
-        display_image(image)
+        _ = display_image(image, matplotlib_figure = False)
     return image
 
 
@@ -221,7 +221,7 @@ def show_image_and_mask(image,
     # Display and return the image.
     image = convert_figure_to_image()
     if not kwargs.get('no_show', False):
-        display_image(image)
+        _ = display_image(image)
     return image
 
 
@@ -262,6 +262,6 @@ def show_semantic_segmentation_truth_and_prediction(image,
     # Display and return the image.
     image = convert_figure_to_image()
     if not kwargs.get('no_show', False):
-        display_image(image)
+        _ = display_image(image)
     return image