log-to-heatmap.py 12 KB

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