def SmartCopyBags(from_dir, to_dir):
"""Copy a task but filter useless sensor data."""
bags = sorted(glob.glob(os.path.join(from_dir, '*.bag')))
if len(bags) == 0:
return
if not os.path.exists(to_dir):
os.makedirs(to_dir)
events = CollectEvents(bags)
next_event = 0
for from_bag in bags:
to_bag = os.path.join(to_dir, os.path.basename(from_bag))
print('Copy bag: {} -> {}'.format(from_bag, to_bag))
# Do the copy
with Bag(from_bag, 'r') as bag_in, Bag(to_bag, 'w') as bag_out:
for topic, msg, t in bag_in.read_messages():
# For small size topics, always copy.
if topic not in K_LARGE_SIZE_TOPICS:
bag_out.write(topic, msg, t)
continue
msg_sec = t.to_sec()
while next_event < len(events) and events[next_event] < msg_sec:
next_event += 1
# For large size topics, only copy when it's near an event.
if (next_event < len(events) and events[next_event] - msg_sec <
K_COPY_LARGE_SIZE_TOPICS_SECONDS_BEFORE_EVENT):
bag_out.write(topic, msg, t)
def bag_to_clock(filename='clock.bag'):
bagfile = Bag(filename)
clock = []
for topic, msg, timestamp in bagfile.read_messages():
if topic == '/clock':
clock.append([timestamp.to_sec(), msg.clock.to_sec()])
return numpy.array(clock)
def bag_to_power(filename='power.bag'):
bagfile = Bag(filename)
energies = {}
power = []
for topic, msg, timestamp in bagfile.read_messages():
if topic.find('energy') != -1:
energies[topic] = msg.data
power.append([timestamp.to_sec(), sum(energies.values())])
bagfile.close()
return numpy.array(power)
def calculate(self, bag_file):
"""Calculate mileage."""
last_pos = None
cur_mode = 'Unknown'
mileage = collections.defaultdict(lambda: 0.0)
with Bag(bag_file, 'r') as bag:
for topic, msg, t in bag.read_messages(
topics=[kChassisTopic, kLocalizationTopic]):
if topic == kChassisTopic:
# Mode changed
if cur_mode != msg.driving_mode:
if (cur_mode == Chassis.COMPLETE_AUTO_DRIVE and
msg.driving_mode == Chassis.EMERGENCY_MODE):
self.disengagements += 1
cur_mode = msg.driving_mode
# Reset start position.
last_pos = None
elif topic == kLocalizationTopic:
cur_pos = msg.pose.position
if last_pos:
# Accumulate mileage, from xyz-distance to miles.
mileage[cur_mode] += 0.000621371 * math.sqrt(
(cur_pos.x - last_pos.x)**2 +
(cur_pos.y - last_pos.y)**2 +
(cur_pos.z - last_pos.z)**2)
last_pos = cur_pos
self.auto_mileage += mileage[Chassis.COMPLETE_AUTO_DRIVE]
self.manual_mileage += (mileage[Chassis.COMPLETE_MANUAL] +
mileage[Chassis.EMERGENCY_MODE])
def rosbaginfo(self, req):
response = RecordingGetBagInfoResponse()
response.filesize = int(os.path.getsize(req.filename) / 1e+6)
bag = Bag(req.filename, 'r')
response.filename = req.filename
response.starttime = bag.get_start_time()
response.endtime = bag.get_end_time()
response.topics = []
topicmap = bag.get_type_and_topic_info()[1]
for topic in topicmap.keys():
topicinfo = BagTopicInfo()
topicinfo.topicname = topic
topicinfo.messagecount = topicmap[topic].message_count
response.topics.append(topicinfo)
return response
def return_fps(filename):
info_dict = yaml.load(Bag(filename, 'r')._get_yaml_info())
dur = info_dict['duration'] #video length in secs
for d in info_dict['topics']:
topic = d['topic']
if d['topic'] == '/device_0/sensor_1/Color_0/image/data':
n_frames = d['messages']
fps = n_frames / dur
print(f'{filename} \ndur = {dur} secs, fps={fps}')
def extract(cls,
automan_info,
file_path,
topics,
output_dir,
raw_data_info,
calibfile=None):
extrinsics_mat, camera_mat, dist_coeff = None, None, None
if calibfile:
try:
calib_path = calibfile
extrinsics_mat, camera_mat, dist_coeff = cls.__parse_calib(
calib_path)
except Exception:
raise UnknownCalibrationFormatError
candidates, topics = cls.__get_candidates(
automan_info, int(raw_data_info['project_id']),
int(raw_data_info['original_id']))
topic_msgs = {}
for topic in topics:
topic_msgs[topic] = ""
try:
count = 0
with Bag(file_path) as bag:
for topic, msg, t in bag.read_messages():
if topic in topics:
topic_msgs[topic] = msg
if all(msg != '' for msg in topic_msgs.values()):
count += 1
for c in candidates:
save_msg = topic_msgs[c['topic_name']]
output_path = output_dir + str(c['candidate_id']) \
+ '_' + str(count).zfill(6)
if (c['msg_type'] == 'sensor_msgs/PointCloud2'):
cls.__process_pcd(save_msg, output_path)
else:
cls.__process_image(save_msg, msg._type,
output_path, camera_mat,
dist_coeff)
for topic in topics:
topic_msgs[topic] = ''
result = {
'file_path': output_dir,
'frame_count': count,
'name': os.path.basename(path), # FIXME
'original_id': int(raw_data_info['original_id']),
'candidates': raw_data_info['candidates'],
}
return result
except Exception as e:
# FIXME
print(traceback.format_exc())
raise (e)
def analyze(file_path):
try:
bag = Bag(file_path)
dataset_candidates = []
for topic_name, info in bag.get_type_and_topic_info().topics.items(
):
if info.msg_type in MSG_DATA_TYPE_MAP.keys():
candidate = {
"analyzed_info": {
"topic_name": topic_name,
"msg_type": info.msg_type,
},
"data_type": MSG_DATA_TYPE_MAP[info.msg_type],
"frame_count": info.message_count
}
dataset_candidates.append(candidate)
return dataset_candidates
except Exception as e:
# FIXME
raise (e)
def analyze(cls, file_path, label_type):
try:
bag = Bag(file_path)
dataset_candidates = []
for topic_name, info in bag.get_type_and_topic_info().topics.items(
):
if info.msg_type in MSG_DATA_TYPE_MAP.keys():
candidate = {
"analyzed_info": {
"topic_name": topic_name,
"msg_type": info.msg_type,
},
"data_type": MSG_DATA_TYPE_MAP[info.msg_type],
"frame_count": info.message_count
}
dataset_candidates.append(candidate)
if cls.__is_label_type_valid(dataset_candidates, label_type):
return dataset_candidates, 'analyzed'
return [], 'invalid'
except Exception as e:
# FIXME
raise (e)
def __init__(self, path, **config):
super(ROSBag, self).__init__(path)
self.topic_img_left = config['topic_img_left'] \
if 'topic_img_left' in config else None
self.topic_img_right = config['topic_img_right'] \
if 'topic_img_right' in config else None
self.topic_imu = config['topic_imu'] \
if 'topic_imu' in config else None
self.topic_temp = config['topic_temp'] \
if 'topic_temp' in config else None
import yaml
from rosbag.bag import Bag
# pylint: disable=protected-access
self._info = yaml.load(Bag(self.path, 'r')._get_yaml_info())
def get_driving_mode(self, bag_file):
"""get driving mode, which is stored in a dict"""
with Bag(bag_file, 'r') as bag:
mode = {}
mode["status"] = 'UNKNOW'
mode["start_time"] = 0.0
mode["end_time"] = 0.0
for topic, msg, _t in bag.read_messages([kChassisTopic]):
t = long(str(_t)) * pow(10, -9)
cur_status = msg.driving_mode
if mode["status"] != cur_status:
if mode["status"] != 'UNKNOW':
self.driving_mode.append(mode)
mode["status"] = cur_status
mode["start_time"] = t
mode["end_time"] = t
self.driving_mode.append(mode)
def process_file(self, bag_file):
"""
Extract information from bag file.
Return True if we are done collecting all information.
"""
try:
with Bag(bag_file, 'r') as bag:
for _, msg, _ in bag.read_messages(topics=kTopics):
if msg.vehicle.name:
self.vehicle_name = msg.vehicle.name.lower()
if msg.vehicle.license.vin:
self.vehicle_vin = msg.vehicle.license.vin
if self.done():
return True
except:
return False
return self.done()
def bag_to_csv(filename):
bagfile = Bag(filename)
topics = init_topics(filename)
start_time = bagfile.get_start_time()
for topic, msg, timestamp in bagfile.read_messages(topics=topics.keys()):
formatter = topics[topic]['formatter']
data_line = formatter((topic,msg,timestamp,start_time))
topics[topic]['file'].write(', '.join(map(str,data_line)) + '\n')
for topic in topics.keys():
topics[topic]['file'].close()
bagfile.close()
def CollectEvents(bags):
"""Collect interested event timestamps."""
print('Collecting events...', end='')
events = []
cur_driving_mode = None
for bag_file in bags:
with Bag(bag_file, 'r') as bag:
for topic, msg, t in bag.read_messages(
topics=[kChassisTopic, kDriveEventTopic]):
# For disengagement, take the message time as event time.
if topic == kChassisTopic:
if (cur_driving_mode == Chassis.COMPLETE_AUTO_DRIVE
and msg.driving_mode == Chassis.EMERGENCY_MODE):
events.append(t.to_sec())
cur_driving_mode = msg.driving_mode
# For DriveEvent, take the header time as event time.
elif topic == kDriveEventTopic:
events.append(msg.header.timestamp_sec)
print('Collected {} events.'.format(len(events)))
return events
def process_bag(self, bag_file):
"""Process a bag and populate task proto."""
print 'Processing bag', bag_file
with Bag(bag_file, 'r') as bag:
if bag.get_message_count() == 0:
return
bag_pb = self.task.bags.add()
bag_pb.name = os.path.basename(bag.filename)
bag_pb.size = bag.size
bag_pb.size = bag.size
bag_pb.version = bag.version
bag_pb.start_time = bag.get_start_time()
bag_pb.end_time = bag.get_end_time()
bag_pb.msg_count = bag.get_message_count()
for topic, info in bag.get_type_and_topic_info().topics.items():
bag_pb.topics[topic].msg_type = info.msg_type
bag_pb.topics[topic].msg_count = info.message_count
if info.frequency:
bag_pb.topics[topic].frequency = info.frequency
# Scan messages.
for topic, msg, t in bag.read_messages(topics=self.kTopics):
self.processors[topic](msg, t)
def bag_to_blobs(filename):
bagfile = Bag(filename)
blobs = []
b_time = None
curr_blobs = [None, None]
poses = []
start_time = bagfile.get_start_time()
for topic, msg, timestamp in bagfile.read_messages():
if topic[-9:] == 'blob_list':
if len(msg.blobs) == 1:
m_time = msg.header.stamp.to_sec()
r_idx = int(topic[-20])
blob = [msg.blobs[0].x, msg.blobs[0].y]
#blobs.append([m_time,r_idx] + blob)
if b_time is None:
b_time = m_time
curr_blobs[r_idx] = blob
elif abs(b_time - m_time) < 0.001:
curr_blobs[r_idx] = blob
blobs.append([b_time] + curr_blobs[0] + curr_blobs[1])
b_time = None
curr_blobs = [None, None]
else:
b_time = m_time
curr_blobs = [None, None]
curr_blobs[r_idx] = blob
elif topic == '/tf':
for tf_msg in msg.transforms:
p_time = tf_msg.header.stamp.to_sec()
base_frame = tf_msg.header.frame_id
child_frame = tf_msg.child_frame_id
pos = tf_msg.transform.translation
pos = [pos.x, pos.y, pos.z]
ori = tf_msg.transform.rotation
ori = [ori.x, ori.y, ori.z, ori.w]
poses.append([p_time, base_frame, child_frame, pos, ori])
bagfile.close()
return blobs, poses
def print_bag_info(self):
info_dict = yaml.load(Bag(self.bag, 'r')._get_yaml_info())
print(info_dict)
# filename = "/home/mayanksati/PycharmProjects/models/AI-Hamid/bags_tracker/2018-11-24-09-25-36/2018-11-24-09-25-38_0.bag" # filename = "/home/mayank_sati/Documents/ob_apollo_build/mayank_apollo/2019-01-15-13-41-55_0.bag" # filename = "/home/mayank_sati/Documents/datsets/1017_PCD_COLLECTION/truck1/2018-10-17-16-48-57_0.bag" filename = "/home/mayank_sati/Documents/datsets/1017_PCD_COLLECTION/ego_station_veh/2018-10-17-14-53-42_0.bag" bag = rosbag.Bag(filename) rosbag_info = rosbag.Bag(filename).read_messages() # for topic, msg, t in rosbag.Bag(filename).read_messages(): # # for topic, msg, t in bag.read_messages(topics=['chatter', 'numbers']): # print msg # bag.close() # kotaro.send(inputs) # ________________________________________________________________ # getting list of all the topic from rosbag info_dict = yaml.load(Bag(filename, 'r')._get_yaml_info()) print(1) # '/apollo/sensor/velodyne64/PointCloud2' # # bag = rosbag_info # # topics = bag.get_type_and_topic_info()[1].keys() # # types = [] # # for i in range(0, len(bag.get_type_and_topic_info()[1].values())): # # types.append(bag.get_type_and_topic_info()[1].values()[i][0]) # # ________________________________________________________________
# -*- coding: utf-8 -*-
# @Time : 2017/10/5 10:48
# @Author : play4fun
# @File : 1.3-Get summary information about a bag.py
# @Software: PyCharm
"""
1.3-Get summary information about a bag.py:
"""
import subprocess, yaml
info_dict = yaml.load(
subprocess.Popen(['rosbag', 'info', '--yaml', 'input.bag'],
stdout=subprocess.PIPE).communicate()[0])
#相同地
import yaml
from rosbag.bag import Bag
info_dict = yaml.load(Bag('input.bag', 'r')._get_yaml_info())
#!/usr/bin/env python
import rospy
from sensor_msgs.msg import PointCloud2
from genpy.rostime import Time
from rosbag.bag import Bag
def callback_rgb(data, time, bag):
print 'received PointCloud2!'
data.header.stamp = time + rospy.Duration.from_sec(1505388514) # shift time stamp to the year 2017 (for kalibr)
bag.write('/ensenso/depth/points', data, t=data.header.stamp)
def rgb_listener(bag):
counter = 0
while(True):
raw_input("press enter to continue...")
counter += 1
time = rospy.Time(counter)
msg_rgb = rospy.wait_for_message('/ensenso/depth/points', PointCloud2)
callback_rgb(msg_rgb, time,bag)
if __name__ == '__main__':
rospy.init_node('rgb_listener', anonymous=True)
try:
bag = Bag('images_for_calibration.bag', 'w', allow_unindexed=True)
rgb_listener(bag)
finally:
bag.close()
parser.add_argument("directory", help="Directory with rosbags")
parser.add_argument("--print", default=False, help="print out the yamls")
parser.add_argument(
"--topic",
default=None,
help="If you want bags with a specific topic, provide it here")
parser.add_argument("--write", default=False)
parser.add_argument("--write_file", default="bags.txt")
args = parser.parse_args()
topic_count = {}
topic_total_msgs = {}
bag_files = glob.glob(args.directory + "*.bag")
for bag in bag_files:
info_dict = yaml.load(Bag(bag, 'r')._get_yaml_info())
if args.print:
pretty_dict(info_dict)
print("-" * 40)
for topic_dict in info_dict["topics"]:
topic = topic_dict["topic"]
if topic not in topic_count:
topic_count[topic] = []
topic_total_msgs[topic] = 0
topic_count[topic].append(bag)
topic_total_msgs[topic] += topic_dict["messages"]
f = None
if args.write:
def get_duration(filename):
info_dict = yaml.load(Bag(filename, 'r')._get_yaml_info())
return info_dict['duration']
def info(self): # type: () -> str
# TODO : Remove this protected function and use the properties provided by the class Bag
return dumps(
load(Bag(self.__path + '/' + self.bag)._get_yaml_info(),
Loader=Loader))
from common import factory
from common import pipeline
import PyMBVCore as Core
import PyJointTools as jt
import yaml
from rosbag.bag import Bag
from utils import detector_utils
from tqdm import tqdm
bridge = CvBridge()
fake_clb = OpenCVCalib2CameraMeta(LoadOpenCVCalib("res/calib_webcam_mshd_vga.json"))
info_dict = yaml.load(Bag(sys.argv[2], 'r')._get_yaml_info())
offset = float(sys.argv[3])
recording_delay = info_dict['start'] + offset
limbSeq = [[0, 1], [0, 5], [0, 9], [0, 13], [0, 17], # palm
[1, 2], [2, 3], [3,4], # thump
[5, 6], [6, 7], [7, 8], # index
[9, 10], [10, 11], [11, 12], # middle
[13, 14], [14, 15], [15, 16], # ring
[17, 18], [18, 19], [19, 20], # pinky
]
final_result = []
def _center(a, b):
def mono_hand_loop(acq, outSize, config, track=False, paused=False):
print("Initialize WACV18 3D Pose estimator (IK)...")
pose_estimator = factory.HandPoseEstimator(config)
print("Initialize MVA19 CVRL Hand pose net...")
estimator = mva19.Estimator(config["model_file"], config["input_layer"], config["output_layer"])
detection_graph, sess = detector_utils.load_inference_graph()
started = True
delay = {True: 0, False: 1}
p2d = bbox = None
smoothing = config.get("smoothing", 0)
boxsize = config["boxsize"]
stride = config["stride"]
peaks_thre = config["peaks_thre"]
print("Entering main Loop.")
for topic, img_msg, ros_time in tqdm(Bag(sys.argv[2]).read_messages()):
if topic != "camera/rgb/image_raw":
continue
try:
bgr = bridge.imgmsg_to_cv2(img_msg, "bgr8")
bgr = cv2.resize(bgr, (640,480), interpolation = cv2.INTER_AREA)
except Exception as e:
print("Failed to grab", e)
break
clb = fake_clb
# compute kp using model initial pose
points2d = pose_estimator.ba.decodeAndProject(pose_estimator.model.init_pose, clb)
oldKp = np.array(points2d).reshape(-1, 2)
score = -1
result_pose = None
# STEP2 detect 2D joints for the detected hand.
if started:
if bbox is None:
bbox = detector_utils.hand_bbox(bgr,detection_graph,sess)
if bbox is None:
if sys.argv[4]=='1':
cv2.imshow("3D Hand Model reprojection",bgr)
cv2.waitKey(1)
to_AC_format(np.zeros((21,3)),ros_time, 1)
continue
else:
dbox = detector_utils.hand_bbox(bgr,detection_graph,sess)
if dbox is not None:
x,y,w,h = bbox
x1,y1,w1,h1 = dbox
if (x1>x+w or x1+w1<x ) or y1+h1<y or y1>y+h:
bbox = dbox
print("updated")
else:
x,y,w,h = dbox
b = max((w,h,224))
b = int(b + b*0.3)
cx = x + w/2
cy = y + h/2
x = cx-b/2
y = cy-b/2
x = max(0,int(x))
y = max(0,int(y))
x = min(x, bgr.shape[1]-b)
y = min(y, bgr.shape[0]-b)
bbox = [x,y,b,b]
x,y,w,h = bbox
crop = bgr[y:y+h,x:x+w]
img, pad = mva19.preprocess(crop, boxsize, stride)
t = time.time()
hm = estimator.predict(img)
est_ms = (time.time() - t)
# use joint tools to recover keypoints
scale = float(boxsize) / float(crop.shape[0])
scale = stride/scale
ocparts = np.zeros_like(hm[...,0])
peaks = jt.FindPeaks(hm[...,:-1], ocparts, peaks_thre, scale, scale)
# convert peaks to hand keypoints
hand = mva19.peaks_to_hand(peaks, x, y)
if hand is not None:
keypoints = hand
mask = keypoints[:, 2] < peaks_thre
keypoints[mask] = [0, 0, 1.0]
if track:
keypoints[mask, :2] = oldKp[mask]
keypoints[:, 2] = keypoints[:, 2]**3
rgbKp = IK.Observations(IK.ObservationType.COLOR, clb, keypoints)
obsVec = IK.ObservationsVector([rgbKp, ])
score, res = pose_estimator.estimate(obsVec)
if track:
result_pose = list(smoothing * np.array(pose_estimator.model.init_pose) + (1.0 - smoothing) * np.array(res))
else:
result_pose = list(res)
# score is the residual, the lower the better, 0 is best
# -1 is failed optimization.
if track:
if -1 < score: #< 150:
pose_estimator.model.init_pose = Core.ParamVector(result_pose)
else:
print("\n===>Reseting init position for IK<===\n")
pose_estimator.model.reset_pose()
bbox = None
if score > -1: # compute result points
p2d = np.array(pose_estimator.ba.decodeAndProject(Core.ParamVector(result_pose), clb)).reshape(-1, 2)
# scale = w/config.boxsize
bbox = mva19.update_bbox(p2d,bgr.shape[1::-1])
p3d = np.array(pose_estimator.ba.decode(Core.ParamVector(result_pose), clb)).reshape(-1,3)
to_AC_format(p3d,ros_time,1)
viz = np.copy(bgr)
if sys.argv[4] == '1' and started and result_pose is not None:
viz = mva19.visualize_3dhand_skeleton(viz, p2d)
pipeline.draw_rect(viz, "Hand", bbox, box_color=(0, 255, 0), text_color=(200, 200, 0))
cv2.putText(viz, 'Hand pose estimation', (20, 20), 0, 0.7, (50, 20, 20), 1, cv2.LINE_AA)
cv2.imshow("3D Hand Model reprojection", viz)
key = cv2.waitKey(delay[paused])
if key & 255 == ord('p'):
paused = not paused
if key & 255 == ord('q'):
break
if key & 255 == ord('r'):
print("\n===>Reseting init position for IK<===\n")
pose_estimator.model.reset_pose()
bbox = None
print("RESETING BBOX",bbox)
def get_info(self):
info_dict = yaml.load(Bag(self.rosbag, 'r')._get_yaml_info())
return info_dict
class RosbagReader():
def __init__(self, input_bag, save_dir):
self.save_dir = save_dir
Path.mkdir(self.save_dir, parents=True, exist_ok=True)
self.bag = Bag(input_bag)
def print_bag_info(self):
info_dict = yaml.load(Bag(self.bag, 'r')._get_yaml_info())
print(info_dict)
def extract_camera_data(self):
image_topic = '/camera_left/color/image_raw'
bag_name = self.bag.filename.strip(".bag").split('/')[-1]
output_dir = os.path.join(self.save_dir, bag_name, 'camera')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
bridge = CvBridge()
count = 0
for topic, msg, t in self.bag.read_messages(topics=[image_topic]):
cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding="passthrough")
cv2.imwrite(os.path.join(output_dir, "frame%06i.png" % count),
cv_img)
print("Wrote image {}".format(count))
count += 1
def extract_lidar_frames(self, extension, topicList=None):
topic_name = '/front_lidar/velodyne_points'
# bag_name = self.bag.filename.strip(".bag").split('/')[-1]
# save_temp_dir = os.path.join(self.save_dir,bag_name)
i = 0
for topic, msg, t in self.bag.read_messages(topic_name):
# for each instance in time that has data for topicName
# parse data from this instance
save_file = self.save_dir / 'frame_{}.{}'.format(i, extension)
i += 1
if topic == '/front_lidar/velodyne_points':
pc_list = [[p[0], p[1], p[2], p[3]]
for p in pc2.read_points(msg,
skip_nans=True,
field_names=("x", "y", "z",
"intensity"))]
# print(np.array(pc_list, dtype=np.float32).shape)
def read_bag(self, save_to_csv=True):
"""
Return dict with all recorded messages with topics as keys
Args:
save_to_csv (bool, optional): Save data to csv files (one individual file per topic) if True
Returns:
dict: containing all published data points per topic
"""
topics = self.readBagTopicList()
messages = {}
max_it_bag = 10
it_bag = 0
# iterate through topic, message, timestamp of published messages in bag
for topic, msg, _ in self.bag.read_messages(topics=topics):
if type(msg).__name__ == '_sensor_msgs__PointCloud2':
points = np.zeros((msg.height * msg.width, 3))
for pp, ii in zip(
pc2.read_points(msg,
skip_nans=True,
field_names=("x", "y", "z")),
range(points.shape[0])):
points[ii, :] = [pp[0], pp[1], pp[2]]
msg = points
if topic not in messages:
messages[topic] = [msg]
else:
messages[topic].append(msg)
it_bag += 1
if it_bag >= max_it_bag:
break # stop -- too many topics
return messages
def readBagTopicList(self):
"""
Read and save the initial topic list from bag
"""
print("[OK] Reading topics in this bag. Can take a while..")
topicList = []
bag_info = yaml.load(self.bag._get_yaml_info(), Loader=yaml.FullLoader)
for info in bag_info['topics']:
topicList.append(info['topic'])
print('{0} topics found'.format(len(topicList)))
return topicList
def __init__(self, input_bag, save_dir):
self.save_dir = save_dir
Path.mkdir(self.save_dir, parents=True, exist_ok=True)
self.bag = Bag(input_bag)
def calculate(self, bag_file):
"""calculate body sensation, it should be after get driving mode"""
with Bag(bag_file, 'r') as bag:
for topic, msg, _t in bag.read_messages([kLocalizationTopic]):
t = long(str(_t)) * pow(10, -9)
self.timestamp = t
diff_bump_time = t - self._last_bump_time
if diff_bump_time <= BUMP_TIME_THRESHOLD:
continue
acc_x = msg.pose.linear_acceleration.x
acc_y = msg.pose.linear_acceleration.y
acc_z = msg.pose.linear_acceleration.z
if abs(
acc_z
) >= SPEED_UP_THRESHOLD_2 and diff_bump_time >= BUMP_TIME_THRESHOLD:
self._bumps_rollback(t)
self._last_bump_time = t
if self._check_status(t):
self.auto_counts["bumps"] += 1
else:
self.manual_counts["bumps"] += 1
else:
if self._speed_down_2_flag:
if acc_y <= SPEED_DOWN_THRESHOLD_4:
self._speed_down_4_flag = 1
continue
if acc_y <= SPEED_DOWN_THRESHOLD_2:
continue
if self._speed_down_4_flag == 1 \
and t - self._last_speed_down_4_time >= ACCELERATE_TIME_THRESHOLD:
self._last_speed_down_4_time = t
if self._check_status(t):
self.auto_counts["speed_down_4"] += 1
else:
self.manual_counts["speed_down_4"] += 1
elif t - self._last_speed_down_2_time >= ACCELERATE_TIME_THRESHOLD:
self._last_speed_down_2_time = t
if self._check_status(t):
self.auto_counts["speed_down_2"] += 1
else:
self.manual_counts["speed_down_2"] += 1
self._speed_down_2_flag = 0
self._speed_down_4_flag = 0
elif acc_y <= SPEED_DOWN_THRESHOLD_2:
self._speed_down_2_flag = 1
if self._speed_up_2_flag:
if acc_y >= SPEED_UP_THRESHOLD_4:
self._speed_up_4_flag = 1
continue
if acc_y >= SPEED_UP_THRESHOLD_2:
continue
if self._speed_up_4_flag == 1 \
and t - self._last_speed_up_4_time >= ACCELERATE_TIME_THRESHOLD:
self._last_speed_up_4_time = t
if self._check_status(t):
self.auto_counts["speed_up_4"] += 1
else:
self.manual_counts["speed_up_4"] += 1
elif t - self._last_speed_up_2_time >= ACCELERATE_TIME_THRESHOLD:
self._last_speed_up_2_time = t
if self._check_status(t):
self.auto_counts["speed_up_2"] += 1
else:
self.manual_counts["speed_up_2"] += 1
self._speed_up_2_flag = 0
self._speed_up_4_flag = 0
elif acc_y >= SPEED_UP_THRESHOLD_2:
self._speed_up_2_flag = 1
if self._turning_flag:
if abs(acc_x) >= SPEED_UP_THRESHOLD_2:
continue
if t - self._last_turning_time >= ACCELERATE_TIME_THRESHOLD:
self._last_turning_time = t
if self._check_status(t):
self.auto_counts["turning"] += 1
else:
self.manual_counts["turning"] += 1
self._turning_flag = 0
elif abs(acc_x) >= SPEED_UP_THRESHOLD_2:
self._turning_flag = 1
def extract(cls,
automan_info,
file_path,
topics,
output_dir,
raw_data_info,
calibfile=None):
extrinsics_mat, camera_mat, dist_coeff = None, None, None
if calibfile:
try:
calib_path = calibfile
extrinsics_mat, camera_mat, dist_coeff = cls.__parse_calib(
calib_path)
except Exception:
raise UnknownCalibrationFormatError
candidates, topics = cls.__get_candidates(
automan_info, int(raw_data_info['project_id']),
int(raw_data_info['original_id']), raw_data_info['records'])
topic_msgs = {}
for topic in topics:
topic_msgs[topic] = ""
try:
frame_time = []
count = 0
with Bag(file_path) as bag:
for topic, msg, t in bag.read_messages():
if topic in topics:
topic_msgs[topic] = msg
if all(msg != '' for msg in topic_msgs.values()):
count += 1
for c in candidates:
save_msg = topic_msgs[c['topic_name']]
output_path = output_dir + str(c['candidate_id']) \
+ '_' + str(count).zfill(6)
if (c['msg_type'] == 'sensor_msgs/PointCloud2'):
cls.__process_pcd(save_msg, output_path)
else:
cls.__process_image(save_msg, c['msg_type'],
output_path, camera_mat,
dist_coeff)
frame_time.append({
'frame_number': count,
'secs': t.secs,
'nsecs': t.nsecs,
})
for topic in topics:
topic_msgs[topic] = ''
name = os.path.basename(path)
if 'name' in raw_data_info and len(raw_data_info['name']) > 0:
name = raw_data_info['name']
result = {
'file_path': output_dir,
'frame_count': count,
'name': name,
'original_id': int(raw_data_info['original_id']),
'candidates': raw_data_info['candidates'],
'frames': frame_time
}
return result
except Exception as e:
print(e)
raise (e)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="save data from bag file")
parser.add_argument("--bag", required=True, help="Path to bag file")
parser.add_argument("--output", default="bag_output", help="directory for data")
parser.add_argument("--blob-topic", dest="blob_topic",
default="/gatlin/objectscreencoords",
help="topic for blobs")
parser.add_argument("--image-topic", dest="image_topic",
default="/camera/rgb/image_rect_color_throttled",
help="topic for images")
parser.add_argument("--factor", default=10**7, help="Factor for decimal digits")
parser.add_argument("--interactive", action="store_true", help="Show all images to mark blocks")
args = parser.parse_args()
bag = Bag(args.bag)
messages = bag.read_messages()
# Convert bag data to correct format
groups = defaultdict(lambda : defaultdict(lambda : defaultdict(list)))
blobs_to_group_later = []
for topic, msg, time in messages:
if topic == args.blob_topic:
data = [msg.x, msg.y]
blobs_to_group_later.append((data, time))
else:
data = np.fromstring(msg.data, np.uint8)
data = np.reshape(data, (msg.height, msg.width, 3))
groups[time.secs][int(time.nsecs / args.factor)][args.image_topic].append(data)
# Match blobs with images
