Repository: aimhubio/aim (language: python)

Issue:
`TensorboardFolderTracker` missing Image import
## 🐛 Bug

Currently the `TensorboardFolderTracker` will crash if there are any images that need to be parsed. This is because currently `aim.Image` is only imported during type checking, however `_process_tb_event` attempts to create an `Image` instance, without access

### To reproduce

- Create a tensorboard events file that contains an image
- Create a run using the `aim.ext.tensorboard_tracker.tracker.Run` runner.
- Observe the console

### Expected behavior

Images are converted and added to the RocksDB database.

### Environment

- Aim Version 3.17.3
- Python version 3.10.5
- pip version 22.0.4
- OS Ubuntu
- Any other relevant information




Discussion / hints:

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 57c867207b..2017eb2edd 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -13,7 +13,8 @@
 
 ### Fixes
 
--  Convert NaNs and Infs in responses to strings (n-gao)
+- Convert NaNs and Infs in responses to strings (n-gao)
+- Import `Image` and `Audio` for `TensorboardFolderTracker` (alansaul)
 
 ## 3.17.4  May 4, 2023
 
diff --git a/aim/ext/tensorboard_tracker/tracker.py b/aim/ext/tensorboard_tracker/tracker.py
index ce00b5e1e0..2d86d3c92a 100644
--- a/aim/ext/tensorboard_tracker/tracker.py
+++ b/aim/ext/tensorboard_tracker/tracker.py
@@ -10,10 +10,8 @@
 import weakref
 import queue
 
-from typing import TYPE_CHECKING, Any
-
-if TYPE_CHECKING:
-    from aim import Audio, Image
+from typing import Any
+from aim import Audio, Image
 
 
 class TensorboardTracker:


Test patch:
diff --git a/tests/README.md b/tests/README.md
index 3b7f693039..8dbe0c6f2e 100644
--- a/tests/README.md
+++ b/tests/README.md
@@ -5,6 +5,7 @@ Be able to test the correctness of the
  - `aim engine`
  - `aim sdk`
  - `aim ql`
+ - `extensions`
 
 ### Folder Structure
 
@@ -16,6 +17,8 @@ tests
     test_*.py
   ql
     test_*.py
+  ext
+    test_*.py
 ```
 
 ## Run
diff --git a/tests/ext/__init__.py b/tests/ext/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/tests/ext/test_tensorboard_tracker.py b/tests/ext/test_tensorboard_tracker.py
new file mode 100644
index 0000000000..545af772e0
--- /dev/null
+++ b/tests/ext/test_tensorboard_tracker.py
@@ -0,0 +1,113 @@
+from queue import Queue
+import numpy as np
+from PIL import ImageChops as PILImageChops
+import tensorflow as tf
+from tensorboard.compat.proto.summary_pb2 import SummaryMetadata, Summary
+from tensorboard.compat.proto.tensor_pb2 import TensorProto
+from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto
+from tensorboard.compat.proto.event_pb2 import Event
+from torch.utils.tensorboard.summary import image, scalar
+
+from aim import Image
+from aim.ext.tensorboard_tracker.tracker import TensorboardFolderTracker
+
+from tests.base import TestBase
+
+
+def images_same_data(image1: Image, image2: Image) -> bool:
+    """
+    Compare two Aim images to see if they contain the same values
+    """
+    image_diff = PILImageChops.difference(
+        image1.to_pil_image(), image2.to_pil_image()
+    )
+    return image_diff.getbbox() is None
+
+
+class TestTensorboardTracker(TestBase):
+
+    def test__process_tb_image_event(self):
+        # Given
+        queue = Queue()
+        tracker = TensorboardFolderTracker(tensorboard_event_folder='dummy', queue=queue)
+        height, width, channels = 5, 4, 3
+        # Note channels is last
+        image_np = np.random.randint(0, 16, (height, width, channels)).astype(dtype=np.uint8)
+        # Create image summary in standard format
+        image_summary = image(tag='test_image', tensor=image_np, dataformats='HWC')
+        event = Event(summary=image_summary)
+
+        # When
+        tracker._process_tb_event(event)
+
+        # Then
+        tracked_image = queue.get().value
+        original_image = Image(image_np)
+        self.assertTrue(isinstance(tracked_image, Image))
+        self.assertTrue(tracked_image.size == original_image.size)
+        self.assertTrue(images_same_data(tracked_image, original_image))
+
+    def test__process_tb_image_plugin_event(self):
+        # Given
+        queue = Queue()
+        tracker = TensorboardFolderTracker(tensorboard_event_folder='dummy', queue=queue)
+        height, width, channels = 5, 4, 3
+        # Note channels is last
+        image_np = np.random.randint(0, 16, (height, width, channels)).astype(dtype=np.uint8)
+        # Create image summary in format of plugin
+        plugin_data = SummaryMetadata.PluginData(plugin_name='images')
+        smd = SummaryMetadata(plugin_data=plugin_data, )
+        tensor = TensorProto(dtype='DT_STRING',
+                             string_val=[
+                                 f"{height}".encode(encoding='utf_8'),
+                                 f"{width}".encode(encoding='utf_8'),
+                                 tf.image.encode_png(image_np).numpy(),
+                             ],
+                             tensor_shape=TensorShapeProto(dim=[TensorShapeProto.Dim(size=3)]))
+
+        image_summary = Summary(
+            value=[Summary.Value(tag='test_image', metadata=smd, tensor=tensor)]
+        )
+        event = Event(summary=image_summary)
+
+        # When
+        tracker._process_tb_event(event)
+
+        # Then
+        tracked_image = queue.get().value
+        original_image = Image(image_np)
+        self.assertTrue(isinstance(tracked_image, Image))
+        self.assertTrue(tracked_image.size == original_image.size)
+        self.assertTrue(images_same_data(tracked_image, original_image))
+
+    def test__process_tb_scalar_simple_value_event(self):
+        # Given
+        queue = Queue()
+        tracker = TensorboardFolderTracker(tensorboard_event_folder='dummy', queue=queue)
+        scalar_np = np.array(0.32, dtype=np.float32)
+        scalar_summary = scalar('test_scalar', scalar_np, new_style=False)
+        event = Event(summary=scalar_summary)
+
+        # When
+        tracker._process_tb_event(event)
+
+        # Then
+        tracked_scalar = queue.get().value
+        self.assertTrue(isinstance(tracked_scalar, float))
+        self.assertTrue(np.allclose(tracked_scalar, scalar_np))
+
+    def test__process_tb_scalar_plugin_event(self):
+        # Given
+        queue = Queue()
+        tracker = TensorboardFolderTracker(tensorboard_event_folder='dummy', queue=queue)
+        scalar_np = np.array(0.32, dtype=np.float32)
+        scalar_summary = scalar('test_scalar', scalar_np, new_style=True)
+        event = Event(summary=scalar_summary)
+
+        # When
+        tracker._process_tb_event(event)
+
+        # Then
+        tracked_scalar = queue.get().value
+        self.assertTrue(isinstance(tracked_scalar, np.ndarray))
+        self.assertTrue(np.allclose(tracked_scalar, scalar_np))