Expose video fps to api (#141)

* Expose video fps to api

* Dummy to force new test

demo.py

@@ -29,7 +29,7 @@
                                              "arctanx": np.arctan(n_iter)}, n_iter)
     x = torch.rand(32, 3, 64, 64) # output from network
     if n_iter%10==0:
-        x = vutils.make_grid(x, normalize=True, scale_each=True)   
+        x = vutils.make_grid(x, normalize=True, scale_each=True)
         writer.add_image('Image', x, n_iter) # Tensor
         #writer.add_image('astronaut', skimage.data.astronaut(), n_iter) # numpy
         #writer.add_image('imread', skimage.io.imread('screenshots/audio.png'), n_iter) # numpy
@@ -42,21 +42,21 @@
         for name, param in resnet18.named_parameters():
             writer.add_histogram(name, param, n_iter)
         writer.add_pr_curve('xoxo', np.random.randint(2, size=100), np.random.rand(100), n_iter) #needs tensorboard 0.4RC or later
-        writer.add_pr_curve_raw('prcurve with raw data', true_positive_counts, 
-            false_positive_counts, 
+        writer.add_pr_curve_raw('prcurve with raw data', true_positive_counts,
+            false_positive_counts,
             true_negative_counts,
             false_negative_counts,
             precision,
             recall, n_iter)
 # export scalar data to JSON for external processing
 writer.export_scalars_to_json("./all_scalars.json")
-            
+
 dataset = datasets.MNIST('mnist', train=False, download=True)
 images = dataset.test_data[:100].float()
 label = dataset.test_labels[:100]
 features = images.view(100, 784)
-writer.add_embedding(features, metadata=label, label_img=images.unsqueeze(1))        
-writer.add_embedding(features, global_step=1, tag='noMetadata')        
+writer.add_embedding(features, metadata=label, label_img=images.unsqueeze(1))
+writer.add_embedding(features, global_step=1, tag='noMetadata')
 dataset = datasets.MNIST('mnist', train=True, download=True)
 images_train = dataset.train_data[:100].float()
 labels_train = dataset.train_labels[:100]
@@ -74,6 +74,8 @@
 # VIDEO
 vid_images = dataset.train_data[:16*48]
 vid = vid_images.view(16, 1, 48, 28, 28)  # BxCxTxHxW
+
 writer.add_video('video', vid_tensor=vid)
+writer.add_video('video_1_fps', vid_tensor=vid, fps=1)
 
 writer.close()

tensorboardX/summary.py

@@ -176,16 +176,16 @@ def make_image(tensor):
                          encoded_image_string=image_string)
 
 
-def video(tag, tensor):
+def video(tag, tensor, fps=4):
     tag = _clean_tag(tag)
     tensor = makenp(tensor, 'VID')
     tensor = tensor.astype(np.float32)
     tensor = (tensor * 255).astype(np.uint8)
-    video = make_video(tensor)
+    video = make_video(tensor, fps)
     return Summary(value=[Summary.Value(tag=tag, image=video)])
 
 
-def make_video(tensor):
+def make_video(tensor, fps):
     try:
         import moviepy.editor as mpy
     except ImportError:
@@ -196,7 +196,7 @@ def make_video(tensor):
     t, h, w, c = tensor.shape
 
     # encode sequence of images into gif string
-    clip = mpy.ImageSequenceClip(list(tensor), fps=4)
+    clip = mpy.ImageSequenceClip(list(tensor), fps=fps)
     with tempfile.NamedTemporaryFile() as f:
         filename = f.name + '.gif'
 

tensorboardX/writer.py

@@ -341,7 +341,7 @@ def add_image(self, tag, img_tensor, global_step=None):
         """
         self.file_writer.add_summary(image(tag, img_tensor), global_step)
 
-    def add_video(self, tag, vid_tensor, global_step=None):
+    def add_video(self, tag, vid_tensor, global_step=None, fps=4):
         """Add video data to summary.
 
         Note that this requires the ``moviepy`` package.
@@ -350,11 +350,12 @@ def add_video(self, tag, vid_tensor, global_step=None):
             tag (string): Data identifier
             vid_tensor (torch.Tensor): Video data
             global_step (int): Global step value to record
+            fps (float or int): Frames per second
         Shape:
             vid_tensor: :math:`(B, C, T, H, W)`.
         """
 
-        self.file_writer.add_summary(video(tag, vid_tensor), global_step)
+        self.file_writer.add_summary(video(tag, vid_tensor, fps), global_step)
 
     def add_audio(self, tag, snd_tensor, global_step=None, sample_rate=44100):
         """Add audio data to summary.