Transferring image data from pupil_server to a client over zmq

client.py

# -*- coding: utf-8 -*-
"""
Created on Thu Apr 10 15:44:47 2014

@author: Daniel Schreij
"""

import zmq
import pygame
from OpenGL.GL import *
import ast
import time
import json
import numpy as np
import math

class PupilListener(object):
	def __init__(self, host_ip, port, scr_width=1280, scr_height=720, fullscreen=False):
		self.scr_width  = scr_width
		self.scr_height = scr_height
		
		# The size of the frame as delivered by pupil
		self.stream_size = (1280,720)		
		
		self.port = port
		self.host_ip = host_ip

		self.__init_network(host_ip, port)
		self.__init_screen(scr_width, scr_height, fullscreen)
		
		
	def __init_network(self, host_ip, port):
		#network setup
		port = str(port)
		context = zmq.Context()
		self.socket = context.socket(zmq.SUB)
		host_addr = host_ip + ":"+port
		print "Connecting to " + host_addr
		self.socket.connect("tcp://" + host_addr)
		#filter by messages by stating string 'STRING'. '' receives all messages
		self.socket.setsockopt(zmq.SUBSCRIBE, '')

		
	def __init_screen(self, scr_width, scr_height, fullscreen):
		pygame.init()
		pygame.display.init()
		flags = pygame.DOUBLEBUF|pygame.OPENGL|pygame.HWSURFACE
		if fullscreen:
			flags = flags | pygame.FULLSCREEN
		pygame.display.set_mode((scr_width, scr_height), flags)
		
		self.__init_GL(scr_width, scr_height)
		self.__texture_setup(scr_width, scr_height)


	def __init_GL(self, scr_width, scr_height):
		glViewport(0, 0, scr_width, scr_height)

		glPushAttrib(GL_ENABLE_BIT)
		glDisable(GL_DEPTH_TEST)
		glDisable(GL_CULL_FACE)

		glDepthFunc(GL_ALWAYS)
		glEnable(GL_BLEND)
		glBlendFunc(GL_SRC_ALPHA, GL_ONE_MINUS_SRC_ALPHA)

		glMatrixMode(GL_PROJECTION)
		glPushMatrix()
		glLoadIdentity()
		glOrtho(0.0,  scr_width, scr_height, 0.0, 0.0, 1.0)

		glMatrixMode(GL_MODELVIEW)
		glPushMatrix()

		glColor4f(1,1,1,1)
		glClearColor(0.0, 0.0, 0.0, 1.0)
		glClearDepth(1.0)	
		
		
	def __texture_setup(self, scr_width, scr_height):
		self.destsize = self.calcScaledRes((scr_width, scr_height), self.stream_size)
		self.vidPos = ((scr_width - self.destsize[0]) / 2, (scr_height - self.destsize[1]) / 2)
		
		# Setup texture in OpenGL to render video to
		glEnable(GL_TEXTURE_2D)
		glMatrixMode(GL_MODELVIEW)
		self.textureNo = glGenTextures(1)
		glBindTexture(GL_TEXTURE_2D, self.textureNo)
		glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
		glTexParameterf(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_NEAREST)
		
		# Fill texture with black to begin with.		
		img = np.zeros([self.stream_size[0],self.stream_size[1],3],dtype=np.uint8)
		img.fill(0) 
		glTexImage2D( GL_TEXTURE_2D, 0, GL_RGB, self.stream_size[0], self.stream_size[1], 0, GL_RGB, GL_UNSIGNED_BYTE, img)
		
		# Create display list which draws the quad on which the texture is rendered
		(x,y) = self.vidPos
		(w,h) = self.destsize
						
		self.frameQuad = glGenLists(1)
		
		glNewList(self.frameQuad, GL_COMPILE)
		glColor4f(1,1,1,1)
		glBegin(GL_QUADS)
		
		glTexCoord2f(1.0, 1.0); glVertex3i(x, y, 0)
		glTexCoord2f(0.0, 1.0); glVertex3i(x+w, y, 0)
		glTexCoord2f(0.0, 0.0); glVertex3i(x+w, y+h, 0)
		glTexCoord2f(1.0, 0.0); glVertex3i(x, y+h, 0)
		
		glEnd()
		glEndList()
		
	
	def draw_circle(self, cx, cy, r, num_segments=25):
		num_segments = 10
		theta = 2 * math.pi / float(num_segments) 
		c = math.cos(theta)  #precalculate the sine and cosine
		s = math.sin(theta)
		
		x = r 	#we start at angle = 0 
		y = 0 
	    		
		glPushMatrix()					
		
		glBegin(GL_TRIANGLE_FAN)
		for i in range(0, num_segments):	
			glColor4f(1.0, 0.0, 0.0, 1.0)
			glVertex2f(x + cx, y + cy) #output vertex 
	        
			#apply the rotation matrix
			t = x
			x = c * x - s * y
			y = s * t + c * y
		glEnd()	
		glPopMatrix()

		
	
	def calcScaledRes(self, screen_res, image_res):
		"""Calculate image size so it fits on screen
		Args
			screen_res (tuple)   -  Display window size/Resolution
			image_res (tuple)    -  Image width and height

		Returns
			tuple - width and height of image scaled to window/screen
		"""
		rs = screen_res[0]/float(screen_res[1])
		ri = image_res[0]/float(image_res[1])

		if rs > ri:
			return (int(image_res[0] * screen_res[1]/image_res[1]), screen_res[1])
		else:
			return (screen_res[0], int(image_res[1]*screen_res[0]/image_res[0]))
		
		
	def listen(self):
		keep_going = True
		render_times = []
		
		while keep_going:
			# Receive pupil info
		
			try:
				pupil_data = self.socket.recv(zmq.NOBLOCK)
				event_data = self.socket.recv(zmq.NOBLOCK)
				
				# receive image metadata
				frame_metadata = self.socket.recv(zmq.NOBLOCK)
				if frame_metadata != "None":
					frame_metadata = json.loads(frame_metadata)		
				
				# Receive image
				frame_img = self.socket.recv(zmq.NOBLOCK)
		
				if pupil_data != "None":
					items = pupil_data.split("\n")
					msg_type = items.pop(0)
					
					items = dict([i.split(':') for i in items[:-1] ])
					
					if msg_type == 'Pupil':
						try:
							coords = ast.literal_eval(items['norm_gaze'])
							x = int(coords[0]*self.scr_width)
							y = int((1.0-coords[1])*self.scr_height)	
						except KeyError:
							x = None
							y = None
					else:
						# process non gaze position events from plugins here
						pass
			except zmq.ZMQError:
				frame_img = "None"				
			
			# Do the drawing
			t1 = time.time()
								
			# Clear The Screen And The Depth Buffer
			glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)			
			
			if frame_img != "None":
				glTexSubImage2D( GL_TEXTURE_2D, 0, 0, 0, self.stream_size[0], self.stream_size[1], GL_RGB, GL_UNSIGNED_BYTE, frame_img[::-1])
			glCallList(self.frameQuad)
			
			try:
				self.draw_circle(x,y,20)
			except:
				#print "Gaze tracker lost"
				pass
			
			# Flip the buffer to show frame to screen
			render_times.append(time.time() - t1)
			pygame.display.flip()			
			
			for e in pygame.event.get():
				if e.type == pygame.QUIT:
					keep_going = False
			pygame.event.pump()   # Prevent freezing of screen while dragging window
		
		print "Mean render time {0} ms".format(1000*np.mean(render_times))
		pygame.quit()
		
if __name__== "__main__":
	listener = PupilListener(host_ip="192.168.137.35", port=5000)
	listener.listen()

pupil_server.py

#### Server side ####

## exerpt of pupil_src/shared_modules/pupil_server.py
def update(self,frame,recent_pupil_positions,events):   
    t_interval = .1    # in seconds

    ## Determine if new frame should be sent to client
    try:
        curr_time = time.time()
        send_image = curr_time - self.last_frame_sent > t_interval            
        if send_image:
            self.last_frame_sent = curr_time
    except:
        # if no last_frame_sent var exists an exception is raised. Initialize everything
        self.last_frame_sent = curr_time
        send_image = True
                                    
    # Check if pupil data should be sent
    if len(recent_pupil_positions):
        msg = "Pupil\n"    
        for p in recent_pupil_positions:                
            for key,value in p.iteritems():
                if key not in self.exclude_list:
                    msg +=key+":"+str(value)+'\n'
        self.socket.send( msg, zmq.SNDMORE )
    else:
        self.socket.send( "None", zmq.SNDMORE )

    # Check if event data should be sent
    if len(events):
        msg = 'Event'+'\n'
        for e in events:                
            for key,value in e.iteritems():
                if key not in self.exclude_list:
                    msg +=key+":"+str(value).replace('\n','')+'\n'
        self.socket.send( msg, zmq.SNDMORE )
    else:
        self.socket.send( "None", zmq.SNDMORE )                                                                

    # Send image data, if applicable
    # ZMQ allows to send a numpy array directly over a socket
    # You do need to send metadata about the image to the client, so it can 
    # reconstruct it there
    if send_image:
        metadata = dict(
            dtype=str(frame.img.dtype),
            shape=frame.img.shape
        )
        self.socket.send_json(metadata, zmq.SNDMORE)
        self.socket.send(frame.img, zmq.SNDMORE, copy=True, track=False)
    else:
        self.socket.send( "None", zmq.SNDMORE )
        self.socket.send( "None" )