visual_IKD/collect_data.py at main · HareshKarnan/visual_IKD · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
"""
Script that records sensor information using rosbag record
"""

import os
import rospy
from sensor_msgs.msg import Joy
import signal
import subprocess

# rgb_topic = '/camera/rgb/image_raw/compressed'
rgb_topic = '/webcam/image_raw/compressed'
depth_topic = '/camera/depth/image_raw/compressed'
odom_topic = '/camera/odom/sample'
gyro_topic = '/camera/gyro/sample'
accel_topic = '/camera/accel/sample'
joystick_topic = '/joystick'
tf_topic = '/tf'
tfstatic_topic = '/tf_static'
vescdrive_topic = '/vesc_drive'


class DataCollector:
    def __init__(self):
        self.setup_subscribers()
        self.button_x, self.button_o = 0, 0
        self.pro = None

    def setup_subscribers(self):
        self.joy_sub = rospy.Subscriber('/joystick', Joy, self.joy_cb)
        print('Subscribers setup successfully')

    def joy_cb(self, data):
        self.button_x = data.buttons[0]
        self.button_o = data.buttons[1]

        if self.button_o == 1 and self.pro == None:
            print('Starting data collection')
            # start data collection
            cmd = 'exec rosbag record ' + \
                rgb_topic + ' '+\
                odom_topic + ' '+\
                gyro_topic + ' '+\
                accel_topic + ' '+\
                joystick_topic + ' '+\
                tf_topic + ' '+\
                tfstatic_topic +' '+\
                vescdrive_topic

            self.pro = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True)
            print('Data collection started')


        if self.button_x == 1:

            if self.pro is None:
                print('Press button O before killing the process with X')

            else:
                # kill the process
                print('Ending data collection')
                # os.kill(os.getpid(self.pro.pid), signal.SIGINT)
                self.pro.kill()
                print('Ended data collection')
                self.pro = None


if __name__ == "__main__":
    rospy.init_node('rosbag_recorder_trigger', anonymous=True)
    recorder = DataCollector()
    rospy.spin()