keyboard_waveform

#!/usr/bin/env python

import os
import time
import json
import subprocess

from queue import Queue
from pathlib import Path

import numpy as np
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw

VID_SIZE = (150, 60)

def read_frames():
	command = [
		'ffmpeg',
		'-f', 'pulse',
		'-i', 'alsa_output.pci-0000_0c_00.4.analog-stereo.monitor',
		'-filter_complex', '[0:a]showwaves=s=150x60:mode=cline:split_channels=1:scale=log:colors=00ffff|cd5c5c:rate=8,format=rgb24[v]',
		'-map', '[v]',
		'-pix_fmt', 'rgb24',
		'-f', 'rawvideo',
		'pipe:1'
	]
	process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

	while True:
		raw_frame = process.stdout.read(VID_SIZE[0] * VID_SIZE[1] * 3)
		if not raw_frame:
			break
		frame = np.frombuffer(raw_frame, dtype='uint8')
		frame = frame.reshape((VID_SIZE[1], VID_SIZE[0], 3))
		yield frame


def scale_points(points, image_width, image_height, frame_width, frame_height):
	scaled_points = []
	for point in points:
		x = point[0] * (frame_width / image_width)
		y = point[1] * (frame_height / image_height)
		scaled_points.append([int(x), int(y)])
	return scaled_points


def int_to_padded_hex(i):
	if not 0 <= i <= 255:
		raise ValueError("Input must be in the range 0-255")

	return hex(i)[2:].zfill(2)


def hex_to_rgb(hex_val):
	hex_val = hex_val.lstrip('#')
	return tuple(int(hex_val[i:i+2], 16) for i in (0, 2, 4))


def main():
	# Parse JSON data from file
	json_file = str(Path.home()) + '/.local/share/keyboard_map/keyboard_labels.json'

	with open(json_file) as f:
		kb_data = json.load(f)

	# Extract shape labels and points
	shapes = kb_data['shapes']
	shape_points = {}

	for shape in shapes:
		if shape['shape_type'] != 'rectangle':
			continue

		shape_points[shape['label']] = shape['points']

	scaled_key_coords = { k: scale_points(v, kb_data['imageWidth'], kb_data['imageHeight'], VID_SIZE[0], VID_SIZE[1]) for k, v in shape_points.items() }

	# Iterate over frames and process
	for frame in read_frames():
		key_color_map = { k: frame[v[0][1]:v[1][1],v[0][0]:v[1][0],:].mean(axis=(0,1)).astype('uint8') for k, v in scaled_key_coords.items() }

		key_color_assignments = [
			f"k { keyname } { ''.join([int_to_padded_hex(i) for i in color]) }" for keyname, color in key_color_map.items()
		]

		if os.getenv('DEBUG', '0') == '1':
			plt.imshow(frame)
			plt.show()

			key_color_display = Image.new(mode='RGB', size=(150, 60), color='white')
			drawer = ImageDraw.Draw(key_color_display)

			breakpoint()

			for key_label, color in key_color_map.items():
				x1, y1 = scaled_key_coords[key_label][0]
				x2, y2 = scaled_key_coords[key_label][1]
				drawer.rectangle([x1, y1, x2, y2], fill=tuple(color))

			key_color_display.show()

		for i in range(0, len(key_color_assignments), 30):
			profile = '\n'.join(key_color_assignments[i:i+30]) +  '\nc\n'
			process = subprocess.Popen(['g810-led', '-pp'], stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
			stdout, stderr = process.communicate(profile.encode())
			stdout, stderr = stdout.decode(), stderr.decode()
			if process.returncode != 0:
				raise Exception()
			time.sleep(.01)


if __name__ == '__main__':
	main()