With testing screen, coordinate mapping, CMRmap colormap

python_scripts/heatmap.py

@@ -2,7 +2,6 @@ import numpy as np
 import os
 import pickle
 import logging
-from scipy.ndimage.filters import gaussian_filter
 from PIL import Image, ImageDraw,ImageTk
 from matplotlib import cm
 import sys
@@ -14,7 +13,7 @@ import time
 import argparse
 import subprocess
 import json
-
+from threading import Thread
 import termios, fcntl, os
 
 class Heatmap:
@@ -31,19 +30,114 @@ class Heatmap:
         ])
 
         self.loadCoordinates()
-        self.windowRoot = Tk.Toplevel()
+        self.windowRoot = Tk.Tk()
+        self.windowRoot.geometry('800x600+4000+0')
+        self.windowRoot.attributes("-fullscreen", True)
         imageWindowSize = tuple(metricsSize)
-        self.windowRoot.geometry('%dx%d+%d+%d' % (imageWindowSize[0],imageWindowSize[1],0,0))
+        # self.windowRoot.geometry('%dx%d+%d+%d' % (imageWindowSize[0],imageWindowSize[1],0,0)) # we go full screen so not needed
-        self.canvas = Tk.Canvas(self.windowRoot,width=imageWindowSize[0],height=imageWindowSize[1])
+        self.canvas = Tk.Canvas(self.windowRoot,width=imageWindowSize[0],height=imageWindowSize[1],bd=0, highlightthickness=0, relief='ridge')
-        self.canvas.pack()
+        self.canvas.pack(fill="both", expand=True)
+
+        self.mapSize = [1200,400]
+        self.mapWindowRoot = Tk.Toplevel(master=self.windowRoot)
+        self.mapWindowRoot.geometry('%dx%d' % tuple(self.mapSize))
+        self.mapCanvas = Tk.Canvas(self.mapWindowRoot,width=self.mapSize[0],height=self.mapSize[1],bd=0, highlightthickness=0, relief='ridge')
+        self.mapCanvas.create_rectangle(20, 50, 22, 350, fill="black", width=0, tags="screenTop")
+        self.mapCanvas.create_rectangle(600, 50, 602, 300, fill="black", width=0, tags="screenSide")
+        # self.mapImage = Image.new('RGB', (500,350))
+        # self.mapDraw = ImageDraw.Draw(self.mapImage)
+        self.mapCanvas.pack(fill="both", expand=True)
+
+        self.windowRoot.bind("<Key>", self.onKeyPress)
+        self.mapWindowRoot.bind("<Key>", self.onKeyPress)
+
+        # self.updateWindow()
+        self.currentTargets = [{
+            u'confidence': 0.983333,
+            u'head_rot': [0.270533, -0.0669274, 0.113554],
+            u'gaze_angle': [0.025313, 0.403179],
+            u'fid': 0,
+            u'head_pos': [73.5302, 26.4475, 399.764],
+            'target': [100,100]
+        }]
+
+        # Have a pre-created blurred circle.
+        circle = Image.open("circle.png").convert("L")
+        self.circle = np.array(circle.resize((circle.size[0]/2,circle.size[1]/2))).astype(float)
+        self.circle = self.circle/np.max(self.circle)
+        self.circleRadius = float(self.circle.shape[0]-1)/2
+
+        # test drawing of metrics:
+        # self.addMetric(5,4,1)
+        # self.addMetric(100,100,1)
+        # self.addMetric(200,200,1)
+        self.addMetric(300,300,2)
+        # self.addMetric(1920,1070,2)
+        self.windowRoot.update()
+        self.windowGeometry = self.windowRoot.geometry().split("+")[0].split("x")
+        self.windowGeometry = (int(self.windowGeometry[0]), int(self.windowGeometry[0]))
+        self.logger.info("Window size: %s", self.windowGeometry)
         self.updateWindow()
 
-    def updateFromJson(self, frame):
+    def onKeyPress(self, event=None, charachter=None):
-        self.logger.info("Received %s", frame)
+        if event is None and charachter is not None:
+            k = charachter
+        else:
+            k = event.char
+
+        print("PRESSED %s" % k)
+        if k.isdigit():
+            print(len(self.currentTargets))
+            if len(self.currentTargets):
+                c = int(k)
+                if c == 1:
+                    self.setCoordinate("tl", self.currentTargets[0])
+                elif c == 2:
+                    self.setCoordinate("tr", self.currentTargets[0])
+                elif c == 3:
+                    self.setCoordinate("bl", self.currentTargets[0])
+                elif c == 4:
+                    self.setCoordinate("br", self.currentTargets[0])
+        else:
+            if k == "q":
+                exit()
+
+
+    def addMetric(self,x,y,confidence):
+        s = self.metrics.shape # y,x
+
+        startX = int( self.circleRadius - x if x <= self.circleRadius else 0)
+        endX = int( s[1] - x + self.circleRadius if x >= (s[1] - self.circleRadius) else self.circle.shape[0])
+        startY = int( self.circleRadius - y if y <= self.circleRadius else 0)
+        endY = int( s[0] - y + self.circleRadius if y >= s[0] - self.circleRadius else self.circle.shape[1])
+
+        mStartX = int(max(0, x-self.circleRadius))
+        mEndX = int(min(s[1], x+self.circleRadius+1))
+        mStartY = int(max(0, y-self.circleRadius))
+        mEndY = int(min(s[0], y+self.circleRadius+1))
+
+        self.logger.debug("Add metric at (%(x)d,%(y)d), circle %(startY)d:%(endY)d,%(startX)d:%(endX)d, matrix %(mStartY)d:%(mEndY)d,%(mStartX)d:%(mEndX)d" % locals())
+        circlePart = self.circle[startY:endY,startX:endX] * confidence
+
+        self.metrics[mStartY:mEndY,mStartX:mEndX] += circlePart
+
+    def keepWindowUpdated(self):
+        while True:
+            self.updateWindow()
+
+    def updateFromJson(self, frame):
+        t1 = time.time()
+        self.logger.info("Received %s", frame)
+        currentTargets = []
 
-        newMetrics = np.zeros((self.metricsSize[1], self.metricsSize[0]))
         for face in frame:
-            # {u'confidence': 0.983333, u'head_rot': [0.270533, -0.0669274, 0.113554], u'gaze_angle': [0.025313, 0.403179], u'fid': 0, u'head_pos': [73.5302, 26.4475, 399.764]}
+            # {
+            #     u'confidence': 0.983333,
+            #     u'head_rot': [0.270533, -0.0669274, 0.113554],
+            #     u'gaze_angle': [0.025313, 0.403179],
+            #     u'fid': 0,
+            #     u'head_pos': [73.5302, 26.4475, 399.764]
+            # }
             x, y = self.getTargetOfFace(face)
             self.logger.debug("Face %d on %s", face['fid'], [x,y])
 
@@ -52,24 +146,71 @@ class Heatmap:
             targetInt = (int(targetPoint[0]), int(targetPoint[1]))
             # check if point fits on screen:
             # if so, measure it
-            if targetInt[0] >= 0 and targetInt[1] >= 0 and targetInt[0] < self.metricsSize[1] and targetInt[1] < self.metricsSize[0]:
+            face['target'] = [targetInt[0], targetInt[1]]
-                newMetrics[targetInt[1],targetInt[0]] += float(face['confidence'])
+            currentTargets.append(face)
+            if targetInt[0] >= 0 and targetInt[1] >= 0 and targetInt[1] < self.metricsSize[1] and targetInt[0] < self.metricsSize[0]:
+                self.addMetric(targetInt[0], targetInt[1], face['confidence'])
+
+        t2 = time.time()
+        self.logger.debug("Update took %fs", t2-t1)
+        self.currentTargets = currentTargets
+        # self.updateWindow()
+
+    def updateMap(self):
+        self.mapCanvas.delete("figure")
+        virtualWidth = 300
+        virtualHeight = 250
+        mmPerPixel = abs(self.coordinates['tl'][0] - self.coordinates['tr'][0])/float(virtualWidth)
+
+        for face in self.currentTargets:
+            dx = face['head_pos'][0]/mmPerPixel #left/right
+            dy = face['head_pos'][1]/mmPerPixel # top/down
+            dz = face['head_pos'][2]/mmPerPixel # front/back
+
+            p1x = int(20 + dz)
+            p1y = int(virtualWidth / 2 + 50 + dx)
+
+            p2x = int(600 + dz)
+            p2y = int(virtualHeight / 2 + 50 + dy)
+            self.mapCanvas.create_oval(p1x-3, p1y-3, p1x+3, p1y+3, fill="red", width=0, tags="figure")
+            self.mapCanvas.create_oval(p2x-3, p2y-3, p2x+3, p2y+3, fill="red", width=0, tags="figure")
+
+            p3x = int(20)
+            p3y = int(float(face['target'][0]) / self.metricsSize[0] * virtualWidth + 50)
+
+            p4x = int(600)
+            p4y = int(float(face['target'][1]) / self.metricsSize[1] * virtualHeight + 50)
+            self.mapCanvas.create_line(p1x,p1y, p3x,p3y, fill="green", tags="figure")
+            self.mapCanvas.create_line(p2x,p2y, p4x,p4y, fill="green", tags="figure")
+
+            self.logger.debug("DRAW FACE TO", face, (p1x,p1y), "AND", (p2x,p2y))
+
 
-        self.metrics = self.metrics + gaussian_filter(newMetrics, sigma = 8)
-        self.updateWindow()
 
     def updateWindow(self):
-        normalisedMetrics = self.metrics / (np.max(self.metrics))
+        t1 = time.time()
+        normalisedMetrics = self.metrics / (max(18,np.max(self.metrics)))
         # convert to colormap, thanks to: https://stackoverflow.com/a/10967471
-        normalisedMetrics = np.uint8(cm.plasma(normalisedMetrics)*255)
+        colormap = cm.plasma
+        colormap = cm.CMRmap
+        normalisedMetrics = np.uint8(colormap(normalisedMetrics)*255)
         image = Image.fromarray(normalisedMetrics)
-        wpercent = (self.metricsSize[0] / float(image.size[0]))
+        self.updateMap()
-        hsize = int((float(image.size[1]) * float(wpercent)))
+        # Too clunky for now:
-        image = image.resize((self.metricsSize[0], hsize))
+        # if len(self.currentTargets) > 0:
+        #     c = ImageDraw.Draw(image)
+        #     for t in self.currentTargets:
+        #         c.ellipse((t[0]-20, t[1]-20, t[0]+20, t[1]+20), fill=(100,100,100,100)))
+        #     del c
+        # wpercent = (self.windowGeometry[0] / float(image.size[0]))
+        # hsize = int((float(image.size[1]) * float(wpercent)))
+        image = image.resize((self.windowGeometry[0], self.windowGeometry[1]))
         tkpi = ImageTk.PhotoImage(image)
         self.canvas.delete("IMG")
-        imagesprite = self.canvas.create_image(500,500,image=tkpi, tags="IMG")
+        imagesprite = self.canvas.create_image(self.windowGeometry[0]/2,self.windowGeometry[1]/2 ,image=tkpi, tags="IMG", anchor="center")
         self.windowRoot.update()
+        t2 = time.time()
+        self.logger.debug("Draw took %fs", t2-t1)
 
     def getTargetOfFace(self, face):
         x = np.arctan(face['gaze_angle'][0])*face['head_pos'][2] + face['head_pos'][0]
@@ -102,17 +243,32 @@ class Heatmap:
 
     def setCoordinate(self, pos, face):
         self.coordinates[pos] = self.getTargetOfFace(face)
+        print(self.hasAllCoordinates(), self.coordinates.values())
         if self.hasAllCoordinates():
+            self.logger.warning("Go create new transform")
             self.saveCoordinates()
             self.updateTransform()
 
+    def runCam(self, openface_exec, device):
+        for output in execute([openface_exec, "-device", str(device), "-cam_width", "1280", "-cam_height", "720"]):
+            try:
+                frame = json.loads(output)
+                self.updateFromJson(frame)
+                try:
+                    c = sys.stdin.read(1)
+                    self.onKeyPress(charachter=c)
+                except IOError: pass
+            except Exception as e:
+                self.logger.warning(str(e))
+                self.logger.warning("received %s", output)
+
 
 
 def main(openface_exec, coordinates_filename, device=0):
 
     logging.basicConfig( format='%(asctime)-15s %(name)s %(levelname)s: %(message)s' )
     logger = logging.getLogger(__name__)
-    logger.setLevel(logging.DEBUG)
+    logger.setLevel(logging.WARNING)
     fd = sys.stdin.fileno()
 
     oldterm = termios.tcgetattr(fd)
@@ -126,28 +282,16 @@ def main(openface_exec, coordinates_filename, device=0):
 
         # metrics matrix
         metricsSize = [1920,1080]
+        metricsSize = [1920/2,1080/2]
+        metricsSize = [800,600]
 
         heatmap = Heatmap(metricsSize, logger, coordinates_filename)
+        # heatmap.runCam(openface_exec, device)
+        thread = Thread(target=heatmap.runCam, args=(openface_exec, device))
+        thread.start()
+        heatmap.keepWindowUpdated()
+
 
-        for output in execute([openface_exec, "-device", str(device), "-cam_width", "1280", "-cam_height", "720"]):
-            try:
-                frame = json.loads(output)
-                heatmap.updateFromJson(frame)
-                try:
-                    c = sys.stdin.read(1)
-                    c = int(c)
-                    if c == 1:
-                        heatmap.setCoordinate("tl", frame[0])
-                    elif c == 2:
-                        heatmap.setCoordinate("tr", frame[0])
-                    elif c == 3:
-                        heatmap.setCoordinate("bl", frame[0])
-                    elif c == 4:
-                        heatmap.setCoordinate("br", frame[0])
-                except IOError: pass
-            except Exception as e:
-                logger.warning(str(e))
-                logger.warning("received %s", output)
     finally:
         termios.tcsetattr(fd, termios.TCSAFLUSH, oldterm)
         fcntl.fcntl(fd, fcntl.F_SETFL, oldflags)