From 0a6225a03c145d76f7699303e39525aefd4b5c74 Mon Sep 17 00:00:00 2001
From: Aaron Schaer <aaschaer@uchicago.edu>
Date: Fri, 2 Aug 2019 15:51:54 -0500
Subject: [PATCH] add prediction video code

---
 .gitignore        |  4 +++-
 prediction_vid.py | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 prediction_vid.py

diff --git a/.gitignore b/.gitignore
index feae5c1..b0d6f7c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,4 @@
 *.pyc
-__pycache__
\ No newline at end of file
+*.mp4
+__pycache__
+frames/
\ No newline at end of file
diff --git a/prediction_vid.py b/prediction_vid.py
new file mode 100644
index 0000000..a28242a
--- /dev/null
+++ b/prediction_vid.py
@@ -0,0 +1,44 @@
+import matplotlib.pyplot as plt
+import imageio
+import numpy
+import os
+
+from conversion import *
+from walk_analysis import paths
+
+
+def create_frames(cors):
+    """
+    For each cordinate pair, output a graph to be used as a frame
+    """
+
+    for i in range(len(cors)):
+        fig = plt.figure()
+        plt.ylim(-50, 50)
+        plt.xlim(-50, 50)
+
+        plt.plot(cors[i][0], cors[i][1], "ro")
+        fig.savefig("frames/{}.png".format(i))
+        plt.close()
+
+
+def create_video():
+    """
+    Merges images in /frames into a 60 fps prediction video
+    """
+    fnames = os.listdir("frames/")
+    fnames.sort(key= lambda a : int(a[:-4]))
+    videowriter = imageio.get_writer('prediction_video.mp4', fps=60)
+    for fname in fnames:
+        videowriter.append_data(plt.imread("frames/{}".format(fname)))
+    videowriter.close()
+
+
+positions = []
+for path in paths:
+    for c in path:
+        positions.append((hex_str_to_single(c[0]),
+                          hex_str_to_single(c[1])))
+
+create_frames(positions)
+create_video()
\ No newline at end of file