From d6e54cf59a0d8d2c6225cab8c56818d450fe1641 Mon Sep 17 00:00:00 2001
From: Ruben van de Ven <git@rubenvandeven.com>
Date: Tue, 5 Feb 2019 20:06:22 +0100
Subject: [PATCH] Save & load metrics when outputdir is set

---
 head_pose.py | 29 +++++++++++++++++++++++++----
 1 file changed, 25 insertions(+), 4 deletions(-)

diff --git a/head_pose.py b/head_pose.py
index 63406e0..55ba52a 100644
--- a/head_pose.py
+++ b/head_pose.py
@@ -88,6 +88,11 @@ c.set(4, 720)
 
 predictor_path = "shape_predictor_68_face_landmarks.dat"
 
+if args.output_dir:
+    lastMetricsFilename = os.path.join(args.output_dir, 'last_metrics.p')
+else:
+    lastMetricsFilename = None
+
 detector = dlib.get_frontal_face_detector()
 predictor = dlib.shape_predictor(predictor_path)
 
@@ -97,7 +102,20 @@ screenDrawCorners = np.array([[10,60], [90, 60], [10, 110], [90, 110]])
 metricsSize = [1920,1080]
 metricsSize = [1280,800]
 dataframe = pd.DataFrame(columns=['x','y'])
-metrics = np.zeros((metricsSize[1], metricsSize[0])) # (y, x)
+
+metrics = None
+if lastMetricsFilename and os.path.isfile(lastMetricsFilename):
+    try:
+        with open(lastMetricsFilename, "rb") as fp:
+            metrics = pickle.load(fp)
+        logger.warn("Loaded metrics from {}".format(lastMetricsFilename))
+    except Exception as e:
+        logger.exception(e)
+
+if metrics is None:
+    metrics = np.zeros((metricsSize[1], metricsSize[0])) # (y, x)
+    logger.warn("New metrics")
+
 screenDrawCorners = np.array([[0,0], [1919,0], [0, 1079], [1919,1079]])
 
 def create_perspective_transform_matrix(src, dst):
@@ -461,12 +479,12 @@ while True:
     tm21 = time.time()
     normalisedMetrics = metrics / (np.max(metrics))
     # convert to colormap, thanks to: https://stackoverflow.com/a/10967471
-    normalisedMetrics = np.uint8(cm.plasma(normalisedMetrics)*255)
+    normalisedMetricsColored = np.uint8(cm.plasma(normalisedMetrics)*255)
     tm22 = time.time()
     logger.debug("Max normalised metrics: %f", np.max(normalisedMetrics))
     # logger.info(normalisedMetrics)
     tm23 = time.time()
-    image = Image.fromarray(normalisedMetrics)
+    image = Image.fromarray(normalisedMetricsColored)
     wpercent = (imageWindowSize[0] / float(image.size[0]))
     hsize = int((float(image.size[1]) * float(wpercent)))
     image = image.resize((imageWindowSize[0], hsize))
@@ -507,9 +525,12 @@ while True:
                     )
                 )
             image.save(filename)
+
+            with open(lastMetricsFilename, 'wb') as fp:
+                pickle.dump( metrics, fp )
+
             logger.debug("Saved frame to {}".format(filename))
             lastSaveTime = now
-        pass
 
     # (optionally) very slowly fade out previous metrics:
     metrics = metrics * .9997