Fix fps bug when creating heatmap frames

heatmappy/video.py

@@ -1,4 +1,6 @@
 from collections import defaultdict
+import json
+import os
 import random
 
 from moviepy.editor import *
@@ -11,21 +13,25 @@ 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
+    def heatmap_on_video(self, base_video, points, heat_fps=15):
+        width, height = base_video.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))
+        return CompositeVideoClip([base_video] + list(heatmap_clips))
+
+    def heatmap_on_video_path(self, video_path, points, heat_fps=15):
+        base = VideoFileClip(video_path)
+        return self.heatmap_on_video(base, points, heat_fps)
 
     @staticmethod
     def _frame_points(pts, fps):
         frames = defaultdict(list)
+        interval = 1000 // fps
         for x, y, t in pts:
-            start = (t // fps) * fps
+            start = (t // interval) * interval
             frames[start].append((x, y))
         return frames
 
@@ -39,24 +45,29 @@ 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))
+                   .set_duration((1000/fps)/1000))
 
 
-if __name__ == '__main__':
+def _example_random_points():
     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'
+    return (rand_point(720, 480, 4000) for _ in range(15000))
 
-    img_heatmapper = Heatmapper(colours='default')
+
+def main():
+    example_vid = os.path.join('assets', 'SampleVideo_720x480_1mb.mp4')
+
+    img_heatmapper = Heatmapper(colours='default', point_strength=0.6)
     video_heatmapper = VideoHeatmapper(img_heatmapper)
 
-    video = heatmap_video = video_heatmapper.heatmap_on_video(
+    heatmap_video = video_heatmapper.heatmap_on_video_path(
         video_path=example_vid,
-        points=example_points
+        points=_example_random_points()
     )
 
-    video.write_videofile('out.mp4', bitrate="5000k", fps=24)
+    heatmap_video.write_videofile('out.mp4', bitrate="5000k", fps=24)
+
+if __name__ == '__main__':
+    main()
 

setup.py

@@ -4,18 +4,18 @@
     required = f.read().splitlines()
 
 setup(
-  name = 'heatmappy',
-  packages = ['heatmappy'],
-  version = '0.1.1',
-  description = 'Draw image heatmaps in python',
-  author = 'Lumen Research',
-  author_email = '[email protected]',
-  url = 'https://github.com/LumenResearch/heatmappy',
-  download_url = 'https://github.com/LumenResearch/heatmappy/tarball/0.1.1',
-  keywords = ['image', 'heatmap', 'heat map'],
-  install_requires=required,
-  classifiers = [
-      'Programming Language :: Python :: 3'
-  ],
-  include_package_data=True,
+    name='heatmappy',
+    packages=['heatmappy'],
+    version='0.2.0',
+    description='Draw image heatmaps in python',
+    author='Lumen Research',
+    author_email='[email protected]',
+    url='https://github.com/LumenResearch/heatmappy',
+    download_url='https://github.com/LumenResearch/heatmappy/tarball/0.1.1',
+    keywords=['image', 'heatmap', 'heat map'],
+    install_requires=required,
+    classifiers=[
+        'Programming Language :: Python :: 3'
+    ],
+    include_package_data=True,
 )