keyboards/ergodox/keymaps/algernon/tools/log-to-heatmap.py

   1 #! /usr/bin/env python3
   2 import json
   3 import os
   4 import sys
   5 import re
   6 import argparse
   7 import time
   8
   9 from math import floor
  10 from os.path import dirname
  11 from subprocess import Popen, PIPE, STDOUT
  12 from blessings import Terminal
  13
  14 class Heatmap(object):
  15     coords = [
  16         [
  17             # Row 0
  18             [ 4,  0], [ 4,  2], [ 2,  0], [ 1,  0], [ 2,  2], [ 3,  0], [ 3,  2],
  19             [ 3,  4], [ 3,  6], [ 2,  4], [ 1,  2], [ 2,  6], [ 4,  4], [ 4,  6],
  20         ],
  21         [
  22             # Row 1
  23             [ 8,  0], [ 8,  2], [ 6,  0], [ 5,  0], [ 6,  2], [ 7,  0], [ 7,  2],
  24             [ 7,  4], [ 7,  6], [ 6,  4], [ 5,  2], [ 6,  6], [ 8,  4], [ 8,  6],
  25         ],
  26         [
  27             # Row 2
  28             [12,  0], [12,  2], [10,  0], [ 9,  0], [10,  2], [11, 0], [     ],
  29             [      ], [11,  2], [10,  4], [ 9,  2], [10,  6], [12, 4], [12, 6],
  30         ],
  31         [
  32             # Row 3
  33             [17,  0], [17,  2], [15,  0], [14,  0], [15,  2], [16,  0], [13,  0],
  34             [13,  2], [16,  2], [15,  4], [14,  2], [15,  6], [17,  4], [17,  6],
  35         ],
  36         [
  37             # Row 4
  38             [20,  0], [20,  2], [19,  0], [18,  0], [19,  2], [], [], [], [],
  39             [19,  4], [18,  2], [19,  6], [20,  4], [20,  6], [], [], [], []
  40         ],
  41         [
  42             # Row 5
  43             [     ], [23,  0], [22,  2], [22,  0], [22,  4], [21,  0], [21,  2],
  44             [24, 0], [24,  2], [25,  0], [25,  4], [25,  2], [26,  0], [      ],
  45         ],
  46     ]
  47
  48     def set_attr_at(self, block, n, attr, fn, val):
  49         blk = self.heatmap[block][n]
  50         if attr in blk:
  51             blk[attr] = fn(blk[attr], val)
  52         else:
  53             blk[attr] = fn(None, val)
  54
  55     def coord(self, col, row):
  56         return self.coords[row][col]
  57
  58     @staticmethod
  59     def set_attr(orig, new):
  60         return new
  61
  62     def set_bg(self, coords, color):
  63         (block, n) = coords
  64         self.set_attr_at(block, n, "c", self.set_attr, color)
  65         #self.set_attr_at(block, n, "g", self.set_attr, False)
  66
  67     def set_tap_info(self, coords, count, cap):
  68         (block, n) = coords
  69         def _set_tap_info(o, _count, _cap):
  70             ns = 4 - o.count ("\n")
  71             return o + "\n" * ns + "%.02f%%" % (float(_count) / float(_cap) * 100)
  72
  73         if not cap:
  74             cap = 1
  75         self.heatmap[block][n + 1] = _set_tap_info (self.heatmap[block][n + 1], count, cap)
  76
  77     @staticmethod
  78     def heatmap_color (v):
  79         colors = [ [0.3, 0.3, 1], [0.3, 1, 0.3], [1, 1, 0.3], [1, 0.3, 0.3]]
  80         fb = 0
  81         if v <= 0:
  82             idx1, idx2 = 0, 0
  83         elif v >= 1:
  84             idx1, idx2 = len(colors) - 1, len(colors) - 1
  85         else:
  86             val = v * (len(colors) - 1)
  87             idx1 = int(floor(val))
  88             idx2 = idx1 + 1
  89             fb = val - float(idx1)
  90
  91         r = (colors[idx2][0] - colors[idx1][0]) * fb + colors[idx1][0]
  92         g = (colors[idx2][1] - colors[idx1][1]) * fb + colors[idx1][1]
  93         b = (colors[idx2][2] - colors[idx1][2]) * fb + colors[idx1][2]
  94
  95         r, g, b = [x * 255 for x in (r, g, b)]
  96         return "#%02x%02x%02x" % (int(r), int(g), int(b))
  97
  98     def __init__(self, layout):
  99         self.log = {}
 100         self.total = 0
 101         self.max_cnt = 0
 102         self.layout = layout
 103
 104     def update_log(self, coords):
 105         (c, r) = coords
 106         if not (c, r) in self.log:
 107             self.log[(c, r)] = 0
 108         self.log[(c, r)] = self.log[(c, r)] + 1
 109         self.total = self.total + 1
 110         if self.max_cnt < self.log[(c, r)]:
 111             self.max_cnt = self.log[(c, r)]
 112
 113     def get_heatmap(self):
 114         with open("%s/heatmap-layout.%s.json" % (dirname(sys.argv[0]), self.layout), "r") as f:
 115             self.heatmap = json.load (f)
 116
 117         ## Reset colors
 118         for row in self.coords:
 119             for coord in row:
 120                 if coord != []:
 121                     self.set_bg (coord, "#d9dae0")
 122
 123         for (c, r) in self.log:
 124             coords = self.coord(c, r)
 125             b, n = coords
 126             cap = self.max_cnt
 127             if cap == 0:
 128                 cap = 1
 129             v = float(self.log[(c, r)]) / cap
 130             self.set_bg (coords, self.heatmap_color (v))
 131             self.set_tap_info (coords, self.log[(c, r)], self.total)
 132         return self.heatmap
 133
 134     def get_stats(self):
 135         usage = [
 136             # left hand
 137             [0, 0, 0, 0, 0],
 138             # right hand
 139             [0, 0, 0, 0, 0]
 140         ]
 141         finger_map = [0, 0, 1, 2, 3, 3, 3, 1, 1, 1, 2, 3, 4, 4]
 142         for (c, r) in self.log:
 143             if r == 5: # thumb cluster
 144                 if c <= 6: # left side
 145                     usage[0][4] = usage[0][4] + self.log[(c, r)]
 146                 else:
 147                     usage[1][0] = usage[1][0] + self.log[(c, r)]
 148             elif r == 4 and (c == 4 or c == 9): # bottom row thumb keys
 149                 if c <= 6: # left side
 150                     usage[0][4] = usage[0][4] + self.log[(c, r)]
 151                 else:
 152                     usage[1][0] = usage[1][0] + self.log[(c, r)]
 153             else:
 154                 fc = c
 155                 hand = 0
 156                 if fc >= 7:
 157                     hand = 1
 158                 fm = finger_map[fc]
 159                 usage[hand][fm] = usage[hand][fm] + self.log[(c, r)]
 160         hand_usage = [0, 0]
 161         for f in usage[0]:
 162             hand_usage[0] = hand_usage[0] + f
 163         for f in usage[1]:
 164             hand_usage[1] = hand_usage[1] + f
 165
 166         total = self.total
 167         if total == 0:
 168             total = 1
 169         stats = {
 170             "total-keys": total,
 171             "hands": {
 172                 "left": {
 173                     "usage": round(float(hand_usage[0]) / total * 100, 2),
 174                     "fingers": {
 175                         "pinky": 0,
 176                         "ring": 0,
 177                         "middle": 0,
 178                         "index": 0,
 179                         "thumb": 0,
 180                     }
 181                 },
 182                 "right": {
 183                     "usage": round(float(hand_usage[1]) / total * 100, 2),
 184                     "fingers": {
 185                         "thumb": 0,
 186                         "index": 0,
 187                         "middle": 0,
 188                         "ring": 0,
 189                         "pinky": 0,
 190                     }
 191                 },
 192             }
 193         }
 194
 195         hmap = ['left', 'right']
 196         fmap = ['pinky', 'ring', 'middle', 'index', 'thumb',
 197                 'thumb', 'index', 'middle', 'ring', 'pinky']
 198         for hand_idx in range(len(usage)):
 199             hand = usage[hand_idx]
 200             for finger_idx in range(len(hand)):
 201                 stats['hands'][hmap[hand_idx]]['fingers'][fmap[finger_idx + hand_idx * 5]] = round(float(hand[finger_idx]) / total * 100, 2)
 202         return stats
 203
 204 def dump_all(out_dir, heatmaps):
 205     stats = {}
 206     t = Terminal()
 207     t.clear()
 208     sys.stdout.write("\x1b[2J\x1b[H")
 209
 210     print ('{t.underline}{outdir}{t.normal}\n'.format(t=t, outdir=out_dir))
 211
 212     keys = list(heatmaps.keys())
 213     keys.sort()
 214
 215     for layer in keys:
 216         if len(heatmaps[layer].log) == 0:
 217             continue
 218
 219         with open ("%s/%s.json" % (out_dir, layer), "w") as f:
 220             json.dump(heatmaps[layer].get_heatmap(), f)
 221         stats[layer] = heatmaps[layer].get_stats()
 222
 223         left = stats[layer]['hands']['left']
 224         right = stats[layer]['hands']['right']
 225
 226         print ('{t.bold}{layer}{t.normal} ({total:,} taps):'.format(t=t, layer=layer,
 227                                                                     total=int(stats[layer]['total-keys'] / 2)))
 228         print (('{t.underline}        | ' + \
 229                 'left ({l[usage]:6.2f}%)  | ' + \
 230                 'right ({r[usage]:6.2f}%) |{t.normal}').format(t=t, l=left, r=right))
 231         print ((' {t.bright_magenta}pinky{t.white}  |     {left[pinky]:6.2f}%     |     {right[pinky]:6.2f}%     |\n' + \
 232                 ' {t.bright_cyan}ring{t.white}   |     {left[ring]:6.2f}%     |     {right[ring]:6.2f}%     |\n' + \
 233                 ' {t.bright_blue}middle{t.white} |     {left[middle]:6.2f}%     |     {right[middle]:6.2f}%     |\n' + \
 234                 ' {t.bright_green}index{t.white}  |     {left[index]:6.2f}%     |     {right[index]:6.2f}%     |\n' + \
 235                 ' {t.bright_red}thumb{t.white}  |     {left[thumb]:6.2f}%     |     {right[thumb]:6.2f}%     |\n' + \
 236                 '').format(left=left['fingers'], right=right['fingers'], t=t))
 237
 238 def process_line(line, heatmaps, opts, stamped_log = None):
 239     m = re.search ('KL: col=(\d+), row=(\d+), pressed=(\d+), layer=(.*)', line)
 240     if not m:
 241         return False
 242     if stamped_log is not None:
 243         if line.startswith("KL:"):
 244             print ("%10.10f %s" % (time.time(), line),
 245                    file = stamped_log, end = '')
 246         else:
 247             print (line,
 248                    file = stamped_log, end = '')
 249         stamped_log.flush()
 250
 251     (c, r, l) = (int(m.group (2)), int(m.group (1)), m.group (4))
 252     if (c, r) not in opts.allowed_keys:
 253         return False
 254
 255     heatmaps[l].update_log ((c, r))
 256
 257     return True
 258
 259 def setup_allowed_keys(opts):
 260     if len(opts.only_key):
 261         incmap={}
 262         for v in opts.only_key:
 263             m = re.search ('(\d+),(\d+)', v)
 264             if not m:
 265                 continue
 266             (c, r) = (int(m.group(1)), int(m.group(2)))
 267             incmap[(c, r)] = True
 268     else:
 269         incmap={}
 270         for r in range(0, 6):
 271             for c in range(0, 14):
 272                 incmap[(c, r)] = True
 273
 274         for v in opts.ignore_key:
 275             m = re.search ('(\d+),(\d+)', v)
 276             if not m:
 277                 continue
 278             (c, r) = (int(m.group(1)), int(m.group(2)))
 279             del(incmap[(c, r)])
 280
 281     return incmap
 282
 283 def main(opts):
 284     heatmaps = {"Dvorak": Heatmap("Dvorak"),
 285                 "ADORE": Heatmap("ADORE")
 286     }
 287     cnt = 0
 288     out_dir = opts.outdir
 289
 290     if not os.path.exists(out_dir):
 291         os.makedirs(out_dir)
 292
 293     opts.allowed_keys = setup_allowed_keys(opts)
 294
 295     if not opts.one_shot:
 296
 297         try:
 298             with open("%s/stamped-log" % out_dir, "r") as f:
 299                 while True:
 300                     line = f.readline()
 301                     if not line:
 302                         break
 303                     if not process_line(line, heatmaps, opts):
 304                         continue
 305         except:
 306             pass
 307
 308         stamped_log = open ("%s/stamped-log" % (out_dir), "a+")
 309     else:
 310         stamped_log = None
 311
 312     while True:
 313         line = sys.stdin.readline()
 314         if not line:
 315             break
 316         if not process_line(line, heatmaps, opts, stamped_log):
 317             continue
 318
 319         cnt = cnt + 1
 320
 321         if opts.dump_interval != -1 and cnt >= opts.dump_interval and not opts.one_shot:
 322             cnt = 0
 323             dump_all(out_dir, heatmaps)
 324
 325     dump_all (out_dir, heatmaps)
 326
 327 if __name__ == "__main__":
 328     parser = argparse.ArgumentParser (description = "keylog to heatmap processor")
 329     parser.add_argument ('outdir', action = 'store',
 330                          help = 'Output directory')
 331     parser.add_argument ('--dump-interval', dest = 'dump_interval', action = 'store', type = int,
 332                          default = 100, help = 'Dump stats and heatmap at every Nth event, -1 for dumping at EOF only')
 333     parser.add_argument ('--ignore-key', dest = 'ignore_key', action = 'append', type = str,
 334                          default = [], help = 'Ignore the key at position (x, y)')
 335     parser.add_argument ('--only-key', dest = 'only_key', action = 'append', type = str,
 336                          default = [], help = 'Only include key at position (x, y)')
 337     parser.add_argument ('--one-shot', dest = 'one_shot', action = 'store_true',
 338                          help = 'Do not load previous data, and do not update it, either.')
 339     args = parser.parse_args()
 340     if len(args.ignore_key) and len(args.only_key):
 341         print ("--ignore-key and --only-key are mutually exclusive, please only use one of them!",
 342                file = sys.stderr)
 343         sys.exit(1)
 344     main(args)