Fixes #471 (#621)

* Fix add_video dataformats

* Fix doc string

tensorboardX/summary.py

@@ -20,7 +20,7 @@
 from .proto.plugin_mesh_pb2 import MeshPluginData
 from .proto import layout_pb2
 from .x2num import make_np
-from .utils import _prepare_video, convert_to_HWC
+from .utils import _prepare_video, convert_to_HWC, convert_to_NTCHW
 
 _INVALID_TAG_CHARACTERS = _re.compile(r'[^-/\w\.]')
 
@@ -340,9 +340,10 @@ def make_image(tensor, rescale=1, rois=None, labels=None):
                          encoded_image_string=image_string)
 
 
-def video(tag, tensor, fps=4):
+def video(tag, tensor, fps=4, dataformats="NTCHW"):
     tag = _clean_tag(tag)
     tensor = make_np(tensor)
+    tensor = convert_to_NTCHW(tensor, input_format=dataformats)
     tensor = _prepare_video(tensor)
     # If user passes in uint8, then we don't need to rescale by 255
     if tensor.dtype != np.uint8:

tensorboardX/utils.py

@@ -90,9 +90,17 @@ def make_grid(I, ncols=8):
             i = i + 1
     return canvas
 
-    # if modality == 'IMG':
-    #     if x.dtype == np.uint8:
-    #         x = x.astype(np.float32) / 255.0
+
+def convert_to_NTCHW(tensor, input_format):
+    assert(len(input_format) == 5), "Only 5D tensor is supported."
+    assert(len(set(input_format)) == len(input_format)), "You can not use the same dimension shordhand twice. \
+        input_format: {}".format(input_format)
+    assert(len(tensor.shape) == len(input_format)), "size of input tensor and input format are different. \
+        tensor shape: {}, input_format: {}".format(tensor.shape, input_format)
+    input_format = input_format.upper()
+    index = [input_format.find(c) for c in 'NTCHW']
+    tensor_NTCHW = tensor.transpose(index)
+    return tensor_NTCHW
 
 
 def convert_to_HWC(tensor, input_format):  # tensor: numpy array

tensorboardX/writer.py

@@ -796,7 +796,8 @@ def add_video(
             vid_tensor: numpy_compatible,
             global_step: Optional[int] = None,
             fps: Optional[Union[int, float]] = 4,
-            walltime: Optional[float] = None):
+            walltime: Optional[float] = None,
+            dataformats: Optional[str] = 'NTCHW'):
         """Add video data to summary.
 
         Note that this requires the ``moviepy`` package.
@@ -807,12 +808,13 @@ def add_video(
             global_step: Global step value to record
             fps: Frames per second
             walltime: Optional override default walltime (time.time()) of event
+            dataformats: Specify different permutation of the video tensor
         Shape:
-            vid_tensor: :math:`(N, T, C, H, W)`. The values should lie in [0, 255] for type
-              `uint8` or [0, 1] for type `float`.
+            vid_tensor: :math:`(N, T, C, H, W)`. The values should lie in [0, 255]
+            for type `uint8` or [0, 1] for type `float`.
         """
         self._get_file_writer().add_summary(
-            video(tag, vid_tensor, fps), global_step, walltime)
+            video(tag, vid_tensor, fps, dataformats=dataformats), global_step, walltime)
 
     def add_audio(
             self,