Add support for description

CHANGELOG.md

@@ -1,3 +1,8 @@
+## neptune-tensorboard 1.1.0
+
+### Changes
+- Add support for `description` argument in `tf.summary.image()` ([#66](https://github.com/neptune-ai/neptune-tensorboard/pull/66))
+
 ## neptune-tensorboard 1.0.0
 
 ### Fixes

src/neptune_tensorboard/integration/tensorflow_integration.py

@@ -43,16 +43,17 @@ def track_scalar(name, data, step=None, description=None, run=None, base_namespa
     run[base_namespace]["scalar"][name].append(data)
 
 
-def track_image(name, data, step=None, run=None, base_namespace=None):
-    # expecting 2 or 3 dimensional tensor. If tensor is 4-dimentional,
-    # as in https://www.tensorflow.org/api_docs/python/tf/summary/image
-    # iterate over first dimension to send all images
-    shape = tf.shape(data)
-    if len(shape) >= 4:
-        for num in range(0, shape[0]):
-            run[base_namespace]["image"][name].append(File.as_image(data[num]))
+def track_image(name, data, step=None, run=None, base_namespace=None, description=None):
+    # If number of images (tf.shape(data)[0]) > 1, append images as FileSeries, else upload as an image.
+    # ref: https://www.tensorflow.org/api_docs/python/tf/summary/image
+    k = tf.shape(data)[0]
+    if k > 1:
+        for num in range(k):
+            run[base_namespace]["image"][name].append(File.as_image(data[num]), description=description)
     else:
-        run[base_namespace]["image"][name] = File.as_image(data)
+        if description:
+            warnings.warn(f"neptune-tensorboard: Uploading single image ({name}). Description will be ignored.")
+        run[base_namespace]["image"][name] = File.as_image(data[0])
 
 
 def track_text(name, data, step=None, description=None, run=None, base_namespace=None):

tests/test_tensorflow.py

@@ -18,7 +18,14 @@ def test_logging():
             tf.summary.scalar("scalar_tensor", data=tf.constant(0.1), step=1)
 
             # Image
-            tf.summary.image("zeros", data=tf.zeros([3, 2, 2, 3]), step=1)
+            tf.summary.image(
+                "single_image", data=tf.random.uniform(shape=[1, 8, 8, 1]), step=1, description="singe_image_desc"
+            )
+
+            # Images
+            tf.summary.image(
+                "image_series", data=tf.random.uniform(shape=[3, 8, 8, 1]), step=1, description="image_series_desc"
+            )
 
             # Text
             tf.summary.text("some_text", data="Hello World!", step=1)
@@ -35,5 +42,6 @@ def fn(x):
 
             assert run.exists("tensorboard")
             assert run.exists("tensorboard/scalar/learning_rate")
-            assert run.exists("tensorboard/image/zeros")
+            assert run.exists("tensorboard/image/single_image")
+            assert run.exists("tensorboard/image/image_series")
             assert run.exists("tensorboard/text/some_text")