def extract_samples(
client_header,
client_messages,
srv_header,
srv_messages,
text_entries=None,
mouse_to_action=default_mouse_to_action,
):
samples = []
client = Client(srv_header)
numpy_screen = client.framebuffer.numpy_screen
numpy_screen.set_paint_cursor(True)
server_deque = collections.deque(srv_messages)
text_deque = None if text_entries is None else collections.deque(
text_entries)
cur_text = ""
for idx, (ts, msg) in enumerate(client_messages):
# update the current text
for text in iterate_earlier(text_deque, ts):
cur_text = text
# apply server messages to the framebuffer
for srv_msg in iterate_earlier(server_deque, ts):
if srv_msg.message_type == rfp_server.RfpServer.MessageType.fb_update:
rects = []
for msg_r in srv_msg.message_body.rects:
rect = client.decode_rectangle(msg_r)
if rect is not None:
rects.append(rect)
update = server_messages.FramebufferUpdate(rects)
numpy_screen.flip()
numpy_screen.apply(update)
numpy_screen.flip()
# pass client action to framebuffer to track cursor position
if msg.message_type == 5: # TODO: enum
event = vnc_event.PointerEvent(msg.message_body.pos_x,
msg.message_body.pos_y,
msg.message_body.button_mask)
numpy_screen.flip()
numpy_screen.apply_action(event)
numpy_screen.flip()
# if button was pressed, record the observation and the event
if msg.message_body.button_mask:
img = crop_image(numpy_screen.peek().copy())
action = mouse_to_action(event)
if action is not None:
obs = img if text_entries is None else (img, cur_text)
samples.append((obs, action))
return samples
def _process_rectangles(self):
if self._remaining_rectangles > 0:
self.expect(self.recv_Rectangle, 12)
else:
framebuffer_update = server_messages.FramebufferUpdate(self._rectangles)
self.numpy_screen.apply(framebuffer_update)
self._remaining_rectangles = None
self._rectangles = None
# Once you're done, send another framebufferUpdateRequest.
#
# Empirically, after a few framebufferUpdateRequest's without
# any changes, Xvnc will hold off on sending another
# framebuffer update until the display data changes.
self.send_FramebufferUpdateRequest(incremental=1)
self.expect(self.recv_ServerToClient, 1)
# apply server messages to the framebuffer
while server_deque:
top_ts = server_deque[0][0]
if top_ts > ts:
break
msg = server_deque.popleft()[1]
# framebuffer update
if msg.message_type == rfp_server.RfpServer.MessageType.fb_update:
rects = []
for msg_r in msg.message_body.rects:
rect = client.decode_rectangle(msg_r)
if rect is not None:
rects.append(rect)
update = server_messages.FramebufferUpdate(rects)
numpy_screen.flip()
numpy_screen.apply(update)
numpy_screen.flip()
# pass client action to framebuffer to track cursor position
if msg.message_type == 5: # TODO: enum
event = vnc_event.PointerEvent(msg.message_body.pos_x,
msg.message_body.pos_y,
msg.message_body.button_mask)
numpy_screen.flip()
numpy_screen.apply_action(event)
numpy_screen.flip()
# if button was pressed, record the image
if msg.message_body.button_mask or last_save is None or (