publish screenshot to Baxter head display

baxter_screen_publisher.py

"""Baxter screenshot publisher
Author: chaonan99
Date: 2018/04/14
"""

from Xlib import display, X
from PIL import Image as PILI  ## Because of name confliction
import numpy as np
import cv2
import argparse

import rospy
import cv_bridge
from sensor_msgs.msg import Image

class ScreenPublisher(object):
  """ScreenPublisher publish screenshot to Baxter head display"""
  def __init__(self):
    super(ScreenPublisher, self).__init__()
    self.W = 1920
    self.H = 1200
    dsp = display.Display()
    self.root = dsp.screen().root
    self.pub = rospy.Publisher('/robot/xdisplay', Image, latch=True, queue_size=1)
    self.bridge = cv_bridge.CvBridge()

  def send_image(self, freq):
    """
    Send the image located at the specified path to the head
    display on Baxter.

    Parameters
    ----------
    freq: float
      publishing frequency
    """
    r = rospy.Rate(freq)
    rospy.loginfo("Start publishing!")
    while not rospy.is_shutdown():
      raw = self.root.get_image(0, 0, self.W, self.H, X.ZPixmap, 0xffffffff)
      img = PILI.frombytes("RGB", (self.W, self.H), raw.data, "raw", "BGRX")

      img_np = np.array(img) #this is the array obtained from conversion
      frame = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)
      frame = cv2.resize(frame, (1024, 600))
      msg = self.bridge.cv2_to_imgmsg(frame, encoding="bgr8")
      self.pub.publish(msg)
      ## Sleep to allow for image to be published.
      r.sleep()

def argument_parse():
  parser = argparse.ArgumentParser(description=main.__doc__)
  parser.add_argument('-f', '--freq', type=float, default=1, help='publishing frequency')
  return parser.parse_args()

def main():
  """Baxter screenshot publisher. Capture the current screenshot at
  some frequency and publish (post) to the Baxter head display.
  """
  args = argument_parse()
  rospy.init_node("screen_to_x", anonymous=True)
  sp = ScreenPublisher()
  sp.send_image(args.freq)

if __name__ == '__main__':
  main()