support list of images for add_images (#434)

tensorboardX/writer.py

@@ -534,7 +534,8 @@ def add_image(self, tag, img_tensor, global_step=None, walltime=None, dataformat
 
     def add_images(self, tag, img_tensor, global_step=None, walltime=None, dataformats='NCHW'):
         """Add batched image data to summary.
-
+        Besides pass 4D tensor, you can also pass a list of tensors of the same size.
+        In this case, the ``dataformats`` should be `CHW` or `HWC`.
         Note that this requires the ``pillow`` package.
 
         Args:
@@ -568,6 +569,20 @@ def add_images(self, tag, img_tensor, global_step=None, walltime=None, dataforma
         """
         if self._check_caffe2_blob(img_tensor):
             img_tensor = workspace.FetchBlob(img_tensor)
+        if isinstance(img_tensor, list):  # a list of tensors in CHW or HWC
+            if dataformats.upper() != 'CHW' and dataformats.upper() != 'HWC':
+                print('A list of image is passed, but the dataformat is neither CHW nor HWC.')
+                print('Nothing is written.')
+                return
+            import torch
+            try:
+                img_tensor = torch.stack(img_tensor, 0)
+            except:
+                import numpy as np
+                img_tensor = np.stack(img_tensor, 0)
+
+            dataformats = 'N' + dataformats
+
         self._get_file_writer().add_summary(
             image(tag, img_tensor, dataformats=dataformats), global_step, walltime)
 

tests/test_writer.py

@@ -48,3 +48,8 @@ def test_writer(self):
                                     recall, n_iter)
             # export scalar data to JSON for external processing
             writer.export_scalars_to_json("./all_scalars.json")
+            imgs = []
+            for i in range(5):
+                imgs.append(np.ones((3, 100, 110)))
+            with SummaryWriter() as w:
+                w.add_images('img_list', imgs, dataformats='CHW')