gist

[FreadMalloc]

#!/usr/bin/env python

-- encoding: utf-8 --

import os
import time
import sys
import cv2
import numpy as np

from ur_msgs.msg import IOStates
from superros.logger import Logger
from superros.comm import RosNode
import superros.transformations as transformations
from visionpylib.cameras import CameraRGB

KEY_SPACE = 32
KEY_Q = 113

#⬢⬢⬢⬢⬢➤ NODE
node = RosNode("save_camera_frames")

node.setupParameter("CAMERA_TOPIC", "/ur5_wrist_cam/image_raw")
node.setupParameter("OUTPUT_PATH", "/tmp/saved_frames_iros2019/temp")
node.setupParameter("CAMERA_CONFIGURATION", 'microsoft_live_camera_focus5.yml') # asus_camera_1_may2017.yml
node.setupParameter("SAVE_WITH_TF", True)
node.setupParameter("TF_BASE", "base_link")
node.setupParameter("TF_TARGET", "wrist_3_link")
node.setupParameter("FIRST_FRAME_INDEX", 0)
node.setupParameter("hz", 60)

node.setHz(1)
tf_enabled = node.getParameter("SAVE_WITH_TF")
camera_calibration_file = node.getParameter("CAMERA_CONFIGURATION")

#⬢⬢⬢⬢⬢➤ Create sCamera Proxy
camera_file = node.getFileInPackage('vision_pkg', 'data/camera_calibration/' + camera_calibration_file)

camera = CameraRGB(
configuration_file=camera_file,
rgb_topic=node.getParameter("CAMERA_TOPIC")
)

frame_save_counter = node.getParameter("FIRST_FRAME_INDEX")

print("@@@@@@@@@@@@@@@@@@ init done @@@@@@@@@@@@@@@@@@@@@")

saving_buffer = 0
saving = False
def cameraCallback(frame):
global frame_save_counter, saving_buffer, saving
""" Camera callback. produce FrameRGBD object """
output = frame.rgb_image.copy()

tf_target = None
if tf_enabled:
    tf_target = node.retrieveTransform(
        node.getParameter("TF_TARGET"),
        node.getParameter("TF_BASE"),
        -1
    )
    if tf_target == None:
        print("no tf target")
        return


if saving_buffer>0:
    if not os.path.exists(node.getParameter("OUTPUT_PATH")):
            os.makedirs(node.getParameter("OUTPUT_PATH"))

    coutner_str = str(frame_save_counter).zfill(5)
    frame_name = "frame_{}.png".format(coutner_str)
    tffile_name = "pose_{}.txt".format(coutner_str)
    frame_save_counter += 1
    filename = os.path.join(
        node.getParameter("OUTPUT_PATH"), frame_name)
    filenametf = os.path.join(
        node.getParameter("OUTPUT_PATH"), tffile_name)
    cv2.imwrite(filename, output)
    if tf_enabled:
        np.savetxt(
            filenametf,
            transformations.KDLtoNumpyVector(tf_target)
        )
    Logger.log("Saved frames:{}".format(coutner_str))
    saving_buffer-=1
    

cv2.imshow("output", output)
c = cv2.waitKey(1)


if c > 0:
    if c == KEY_SPACE:
        saving_buffer=10
    if c == KEY_Q:
        print("EXIT")
        sys.exit(0)

#⬢⬢⬢⬢⬢➤ Camera Msgs Callback
camera.registerUserCallabck(cameraCallback)

def iocallback(msg):
global frame_save_counter, saving_buffer, saving
for dout in msg.digital_out_states:
if dout.pin==16 and dout.state == True:
if not saving:
saving = True
saving_buffer = 10
if dout.pin==16 and dout.state == False:
if saving:
saving = False

node.createSubscriber("/ur_driver/io_states",IOStates,iocallback)

while node.isActive():
node.tick()