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