Can create video heatmaps

heatmappy/__init__.py

@@ -1,4 +1,6 @@
 from heatmappy.heatmap import Heatmapper,\
                               GreyHeatMapper,\
                               PILGreyHeatmapper,\
-                              PySideGreyHeatmapper
+                              PySideGreyHeatmapper
+
+from heatmappy.video import VideoHeatmapper

heatmappy/heatmap.py

@@ -89,7 +89,6 @@ def _colourised(self, img):
         rgba_img = self.cmap(arr, bytes=True)
         return Image.fromarray(rgba_img)
 
-
     @staticmethod
     def _cmap_from_image_path(img_path):
         img = Image.open(img_path)

heatmappy/video.py

@@ -0,0 +1,62 @@
+from collections import defaultdict
+import random
+
+from moviepy.editor import *
+import numpy as np
+
+from heatmappy import Heatmapper
+
+
+class VideoHeatmapper:
+    def __init__(self, heatmapper):
+        self.heatmapper = heatmapper
+
+    def heatmap_on_video(self, video_path, points, heat_fps=25):
+        base = VideoFileClip(video_path)
+        width, height = base.size
+
+        frame_points = self._frame_points(points, heat_fps)
+        heatmap_frames = self._heatmap_frames(width, height, frame_points)
+        heatmap_clips = self._heatmap_clips(heatmap_frames, heat_fps)
+
+        return CompositeVideoClip([base] + list(heatmap_clips))
+
+    @staticmethod
+    def _frame_points(pts, fps):
+        frames = defaultdict(list)
+        for x, y, t in pts:
+            start = (t // fps) * fps
+            frames[start].append((x, y))
+        return frames
+
+    def _heatmap_frames(self, width, height, frame_points):
+        for frame_start, points in frame_points.items():
+            heatmap = self.heatmapper.heatmap(width, height, points)
+            yield frame_start, np.array(heatmap)
+
+    @staticmethod
+    def _heatmap_clips(heatmap_frames, fps):
+        for frame_start, heat in heatmap_frames:
+            yield (ImageClip(heat)
+                   .set_start(frame_start/1000)
+                   .set_duration(fps/1000)
+                   .set_fps(fps))
+
+
+if __name__ == '__main__':
+    def rand_point(max_x, max_y, max_t):
+        return random.randint(0, max_x), random.randint(0, max_y), random.randint(0, max_t)
+
+    example_points = (rand_point(720, 480, 4000) for _ in range(15000))
+    example_vid = 'assets\SampleVideo_720x480_1mb.mp4'
+
+    img_heatmapper = Heatmapper(colours='default')
+    video_heatmapper = VideoHeatmapper(img_heatmapper)
+
+    video = heatmap_video = video_heatmapper.heatmap_on_video(
+        video_path=example_vid,
+        points=example_points
+    )
+
+    video.write_videofile('out.mp4', bitrate="5000k", fps=24)
+

requirements.txt

@@ -1,5 +1,6 @@
 cycler==0.9.0
 matplotlib==1.5.0
+moviepy==0.2.2.11
 numpy==1.10.1
 Pillow==3.0.0
 pyparsing==2.0.6