log-to-heatmap.py 12 KB

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