Host: Made rendering pipeline a bit more flexible

1 files changed, 25 insertions, 17 deletions

diff --git a/host/matelight.py b/host/matelight.py
index 9966aee..886e0f1 100644
--- a/host/matelight.py
+++ b/host/matelight.py
@@ -3,28 +3,36 @@ from itertools import product
 from ctypes import c_size_t, c_uint8, c_void_p, c_float, CDLL, Structure, POINTER
 import numpy as np
 import time
+import atexit
 
-from config import *
+import config
 
 ml = CDLL('./libml.so')
 ml.matelight_open.restype = c_void_p
 
 if ml.matelight_usb_init():
 	raise OSError('Cannot initialize USB library')
-matelights = ml.matelight_open()
-if matelights is None:
-	raise ImportError('Cannot open any Mate Light devices')
-
-dbuf = np.zeros(DISPLAY_WIDTH*DISPLAY_HEIGHT*4, dtype=np.uint8)
-def sendframe(framedata):
-	""" Send a frame to the display
-
-	The argument contains a h * w array of 3-tuples of (r, g, b)-data or 4-tuples of (r, g, b, a)-data where the a
-	channel is ignored.
-	"""
-	# just use the first Mate Light available
-	rgba = len(framedata) == DISPLAY_WIDTH*DISPLAY_HEIGHT*4
-	global dbuf
-	np.copyto(dbuf[:640*(3+rgba)], np.frombuffer(framedata, dtype=np.uint8))
-	ml.matelight_send_frame(matelights, dbuf.ctypes.data_as(POINTER(c_uint8)), c_size_t(CRATES_X), c_size_t(CRATES_Y), c_float(BRIGHTNESS), rgba)
+
+atexit.register(ml.matelight_usb_destroy)
+
+class Matelight:
+	def __init__(self, match_serial=None):
+		""" Open the matelight matching the USB serial number given as a bytes object. If match_serial is None, open the
+		first matelight """
+		self.handle = ml.matelight_open(match_serial)
+		self.dbuf = np.zeros(config.frame_size*4, dtype=np.uint8)
+		if self.handle is None:
+			raise ValueError('Cannot find requested matelight.')
+
+	def sendframe(self, framedata):
+		""" Send a frame to the display
+
+		The argument contains a h * w array of 3-tuples of (r, g, b)-data or 4-tuples of (r, g, b, a)-data where the a
+		channel is ignored.
+		"""
+		# just use the first Mate Light available
+		rgba = len(framedata) == config.frame_size*4
+		np.copyto(self.dbuf[:640*(3+rgba)], np.frombuffer(framedata, dtype=np.uint8))
+		ml.matelight_send_frame(self.handle, self.dbuf.ctypes.data_as(POINTER(c_uint8)), c_size_t(config.crates_x),
+				c_size_t(config.crates_y), c_float(config.brightness), rgba)