iosurface.py

iosurface.py

from objc_util import *
from ctypes import *
from contextlib import contextmanager
import ui
import numpy as np
''' define ctypes signatures'''
IOSurfaceCreate=c.IOSurfaceCreate
IOSurfaceCreate.argtypes=[c_void_p]
IOSurfaceCreate.restype=c_void_p

IOSurfaceGetBaseAddress=c.IOSurfaceGetBaseAddress
IOSurfaceGetBaseAddress.argtypes=[c_void_p]
IOSurfaceGetBaseAddress.restype=c_void_p

IOSurfaceGetBytesPerRow=c.IOSurfaceGetBytesPerRow
IOSurfaceGetBytesPerRow.argtypes=[c_void_p]
IOSurfaceGetBytesPerRow.restype=c_size_t

IOSurfaceGetPlaneCount=c.IOSurfaceGetPlaneCount
IOSurfaceGetPlaneCount.argtypes=[c_void_p]
IOSurfaceGetPlaneCount.restype=c_size_t

IOSurfaceLock=c.IOSurfaceLock
IOSurfaceLock.argtypes=[c_void_p, c_uint32, POINTER(c_uint32) ]
IOSurfaceLock.restype=c_int

IOSurfaceUnlock=c.IOSurfaceUnlock
IOSurfaceUnlock.argtypes=[c_void_p, c_uint32, POINTER(c_uint32) ]
IOSurfaceUnlock.restype=c_int

IOSurfaceGetPixelFormat=c.IOSurfaceGetPixelFormat
IOSurfaceGetPixelFormat.argtypes=[c_void_p]
IOSurfaceGetPixelFormat.restype=c_int32



kCVPixelFormat_32RGBA=int.from_bytes(b'RGBA', byteorder='big')

class IOSurfaceWrapper(object):
	'''Wraps an IOSurface, exposing a numpy array and view.
			.array=numpy array (hxwx4 channels)
			.view = ui.View with display layer boind to .array contents
			.Lock() context manager for updating array. use .Lock(True) to update, and .Lock(False) to delay redraw, with .redraw() to manually force render
			
	  s.Lock context manager must wrap all updates to array. 
	  i.e:
	  # redraw when context manager exits:
	  with s.Lock(redraw=True):
	  	   s.arrray[...] #manipulate data
	 
	 #delayed redraw, 
	  with s.Lock(redraw=False):
	  	   s.arrray[...] #manipulate data
	  with s.Lock(redraw=False):
	  	   s.arrray[...] #manipulate data
	  s.redraw()
	  
	'''
	def __init__(self, width=1024, height=768):
		bpp=4 #bytes per pixel, one per color plus alpha
		properties=ns(
			{ObjCInstance(c_void_p.in_dll(c,'kIOSurfaceWidth')):width,
			ObjCInstance(c_void_p.in_dll(c,'kIOSurfaceHeight')):height,
			ObjCInstance(c_void_p.in_dll(c,'kIOSurfaceBytesPerElement')):bpp,
			ObjCInstance(c_void_p.in_dll(c, 'kIOSurfacePixelFormat')) : 															kCVPixelFormat_32RGBA
			})
		self.height=height
		self.width=width
		self.surf=IOSurfaceCreate(ns(properties))
		self.planecount = IOSurfaceGetPlaneCount(self.surf);
		self.base = IOSurfaceGetBaseAddress(self.surf);
		self.stride = IOSurfaceGetBytesPerRow(self.surf);
		
		stridewidth=self.stride//bpp
		
		data=cast(self.base, POINTER(c_uint8*bpp*(stridewidth)*(self.height)) ).contents

		a=np.ctypeslib.as_array(data)
		#handle the 'extra' columns. 
		self.array=a[0:self.height,0:self.width,:]
		self.array[:,:,:]=255 #set to white, opaque
		self.setupview()

	@contextmanager
	def Lock(self, redraw = True):
		try:
			IOSurfaceLock(self.surf, 0, None)
			self.base = IOSurfaceGetBaseAddress(self.surf);
			yield
		finally:
			IOSurfaceUnlock(self.surf, 0, None)
			if redraw:
				self.redraw()
				
	def setupview(self):
		self.view=ui.View(frame=(0,0,self.width,self.height))
		self._layer=self.view.objc_instance.layer()
		self._layer.contentsOpaque=0 #1 may be faster, if alpha channel not needed
		self._layer.contentsFlipped=0 #set to 1 if needed
		#bind IOSurface to the layer
		self._layer.setContents_(c_void_p(self.surf))
		
	@on_main_thread
	def redraw(self):
		self._layer.setContentsChanged()
	def __del__(self):
		self._layer.contents=None
		del self.array
		from objc_util import c, c_void_p
		CFRelease=c.CFRelease
		CFRelease.argtypes=[c_void_p]
		CFRelease.restype=None
		CFRelease(self.surf)
		
if __name__=='__main__':
	s=IOSurfaceWrapper(256,256)

	v=ui.View(frame=(0,0,s.width,s.height))
	iv=ui.ImageView(frame=(0,0,s.width,s.height))
	iv.image=ui.Image.named('test:Lenna')
	v.add_subview(iv)
	v.add_subview(s.view)
	v.present('sheet')


	#pre-allocate, because rand is slow
	I=np.random.randint(255,size=(50,50,4))	
	N=200 #numframes
	import time
	t=time.perf_counter()
	for i in range(N):
			with s.Lock():
				#blit I into place
				x=np.random.randint(s.width-50)
				y=np.random.randint(s.height-50)
				s.array[y:(y+50),x:(x+50),0:4]=I
	
	print('fps:',N/(time.perf_counter()-t))