|
| 1 | +import image_processing as imp |
| 2 | + |
| 3 | +import skylight as skl |
| 4 | +import numpy as np |
| 5 | +import rosbag |
| 6 | + |
| 7 | +import shutil |
| 8 | +import os |
| 9 | +import re |
| 10 | + |
| 11 | +S2I = ["135", "045", "090", "000"] |
| 12 | + |
| 13 | +response_dist_base = os.path.abspath(os.path.join(os.getcwd(), '..', 'sensor_response_distribution')) |
| 14 | + |
| 15 | + |
| 16 | +def rosbag_to_dict(bag): |
| 17 | + """ |
| 18 | + Extract the relevant messages from a rosbag file and store them in a |
| 19 | + dictionary. The dictionary has one key per topic, and each topic contains |
| 20 | + all message data recorded on that topic. Pol op preferred angles are also |
| 21 | + included for completeness. |
| 22 | +
|
| 23 | + :param bag: a rosbag (note not a filename, the actual bag object) |
| 24 | + :returns: a dictionary containing the recorded data |
| 25 | + """ |
| 26 | + # Simplistic but works |
| 27 | + po = [[],[],[],[],[],[],[],[]] |
| 28 | + yaw = [] |
| 29 | + |
| 30 | + # For pol-ops, topics are pol_op_X where x is from 0 to 7 |
| 31 | + topics_of_interest = ["pol_op_{}".format(x) for x in range(8)] |
| 32 | + topics_of_interest.append("yaw") |
| 33 | + |
| 34 | + # These recordings also contain (blurred) image data (topic 'frames') |
| 35 | + # and full IMU data (topic 'odom'). |
| 36 | + for topic,msg,t in bag.read_messages(topics=topics_of_interest): |
| 37 | + if "pol_op_" in topic: |
| 38 | + idx = int(topic.split("_")[2]) |
| 39 | + po[idx].append(msg.data) |
| 40 | + elif "yaw" in topic: |
| 41 | + yaw.append(msg.data) |
| 42 | + |
| 43 | + # Using this construction: |
| 44 | + # po[x] -> recorded data for pol-op unit x |
| 45 | + # po[x][y] -> pol-op x at time y |
| 46 | + # po[x][y][z] -> pol-op x, time y, photodiode z |
| 47 | + # Absolute value is taken as some photodiodes return negative values. Using |
| 48 | + # the absolute values gives the correct sensor output. |
| 49 | + po = np.abs(po).tolist() |
| 50 | + |
| 51 | + # Azimuths for each pol-op unit. po_az[i] corresponds |
| 52 | + # to po[i]. |
| 53 | + po_az = np.radians([0, 90, 180, 270, 45, 135, 225, 315]).tolist() |
| 54 | + |
| 55 | + data_dictionary = dict() |
| 56 | + for idx in range(8): |
| 57 | + data_dictionary["pol_op_{}".format(idx)] = list(po[idx]) |
| 58 | + |
| 59 | + data_dictionary["azimuths"] = po_az |
| 60 | + data_dictionary["yaw"] = yaw |
| 61 | + |
| 62 | + return data_dictionary |
| 63 | + |
| 64 | + |
| 65 | +def read_bagfile(session, images_dir): |
| 66 | + # timestamp = datetime.strptime(session.split(os.path.sep)[-1], '%H-%M__%A_%d-%m-%y') |
| 67 | + # category = session.split(os.path.sep)[-3].replace("_data", "") |
| 68 | + collection = session.split(os.path.sep)[-4].replace("_data", "") |
| 69 | + |
| 70 | + os.chdir(session) |
| 71 | + |
| 72 | + image_file = [x for x in os.listdir() if ".jpg" in x] |
| 73 | + |
| 74 | + if len(image_file) == 0: |
| 75 | + print("Warning: missing sky image file from the session directory.") |
| 76 | + image_file = None |
| 77 | + elif len(image_file) > 1: |
| 78 | + print("Warning: multiple image files in directory.") |
| 79 | + else: |
| 80 | + image_file = image_file[0] |
| 81 | + |
| 82 | + if image_file is not None: |
| 83 | + shutil.copy2(image_file, images_dir) |
| 84 | + |
| 85 | + recordings = [x for x in os.listdir() if ".bag" in x] |
| 86 | + |
| 87 | + if len(recordings) == 0: |
| 88 | + print("There are no rosbag files in the specified directory.") |
| 89 | + |
| 90 | + full_data = dict() |
| 91 | + for rec in recordings: |
| 92 | + bag = rosbag.Bag(rec) |
| 93 | + rec_data = rosbag_to_dict(bag) |
| 94 | + full_data[rec] = rec_data |
| 95 | + |
| 96 | + data_frame = {"IMU": [], "sun_azimuth": [], "sun_elevation": [], |
| 97 | + "I135": [], "I045": [], "I090": [], "I000": [], "device": [], "rotation": [], |
| 98 | + "location": [], "timestamp": [], "longitude": [], "latitude": [], |
| 99 | + "collection": [], "session": []} |
| 100 | + |
| 101 | + image_name = image_file.replace('.jpg', '') |
| 102 | + obs = imp.get_observer_from_file_name(image_name) |
| 103 | + sun_predict = imp.extract_sun_vector_from_name(image_name) |
| 104 | + sun_model = skl.Sun(obs) |
| 105 | + |
| 106 | + for c, key in enumerate(full_data.keys()): |
| 107 | + data = full_data[key] |
| 108 | + |
| 109 | + pol_op_keys = ["pol_op_{}".format(x) for x in range(8)] |
| 110 | + po = [] |
| 111 | + for k in pol_op_keys: |
| 112 | + po.append(data[k]) |
| 113 | + po = np.array(po) |
| 114 | + |
| 115 | + imu = np.array(data["yaw"]) |
| 116 | + imu_drift = np.angle(np.squeeze(sun_predict)) - sun_model.az |
| 117 | + imu = (imu - imu_drift + np.pi) % (2 * np.pi) - np.pi |
| 118 | + |
| 119 | + imus = np.linspace(imu[:-1], imu[1:], 8, endpoint=False) |
| 120 | + yaws = (imus.T - data['azimuths']).T |
| 121 | + |
| 122 | + # # unit 5 photoreceptor 1 is faulty |
| 123 | + # # unit 7 is faulty, so it is ignored |
| 124 | + # i_valid = [0, 1, 2, 3, 4, 5, 6] |
| 125 | + # imu = yaws[i_valid] # .reshape((-1,)) |
| 126 | + # pol = po[i_valid, 1:] # .reshape((-1, 4)) |
| 127 | + imu = yaws |
| 128 | + pol = po[:, 1:] |
| 129 | + |
| 130 | + sun_azi = np.rad2deg((sun_model.az + np.pi) % (2 * np.pi) - np.pi) |
| 131 | + sun_ele = np.rad2deg((sun_model.alt + np.pi) % (2 * np.pi) - np.pi) |
| 132 | + for unit in range(imu.shape[0]): |
| 133 | + data_frame["IMU"].extend(np.rad2deg(imu[unit])) |
| 134 | + data_frame["sun_azimuth"].extend([sun_azi] * imu.shape[1]) |
| 135 | + data_frame["sun_elevation"].extend([sun_ele] * imu.shape[1]) |
| 136 | + for i in range(4): |
| 137 | + data_frame[f"I{S2I[i]}"].extend(pol[unit, :, i]) |
| 138 | + data_frame["device"].extend([unit] * imu.shape[1]) |
| 139 | + data_frame["rotation"].extend([key.split('.')[0]] * imu.shape[1]) |
| 140 | + # data_frame["category"].extend([category] * imu.shape[1]) |
| 141 | + data_frame["location"].extend([obs.city] * imu.shape[1]) |
| 142 | + data_frame["timestamp"].extend([obs.date] * imu.shape[1]) |
| 143 | + data_frame["longitude"].extend([obs.lon] * imu.shape[1]) |
| 144 | + data_frame["latitude"].extend([obs.lat] * imu.shape[1]) |
| 145 | + data_frame["collection"].extend([collection] * imu.shape[1]) |
| 146 | + data_frame["session"].extend([image_file] * imu.shape[1]) |
| 147 | + |
| 148 | + return data_frame |
| 149 | + |
| 150 | + |
| 151 | +def read_pd_bagfile(session=None): |
| 152 | + if session is None: |
| 153 | + session = response_dist_base |
| 154 | + |
| 155 | + os.chdir(session) |
| 156 | + map_file = [x for x in os.listdir() if ".txt" in x][0] |
| 157 | + recordings = [x for x in os.listdir() if ".bag" in x] |
| 158 | + |
| 159 | + full_data = {} |
| 160 | + |
| 161 | + for r in recordings: |
| 162 | + bag = rosbag.Bag(r) |
| 163 | + data_name = r.split('.')[0] |
| 164 | + full_data[data_name] = { |
| 165 | + 'response': [], |
| 166 | + 'time': [] |
| 167 | + } |
| 168 | + for topic, msg, t in bag.read_messages(): |
| 169 | + full_data[data_name]["response"].append(msg.data[2]) |
| 170 | + full_data[data_name]["time"].append(t.to_time()) |
| 171 | + |
| 172 | + full_data["map"] = {} |
| 173 | + with open(map_file, 'r') as f: |
| 174 | + unit_no = -1 |
| 175 | + for _ in range(42): |
| 176 | + line = f.readline() |
| 177 | + if 'PD' in line: |
| 178 | + result = re.match(r"= PD ([0-9]+) =", line) |
| 179 | + if result is None: |
| 180 | + continue |
| 181 | + unit_no = int(result.group(1)) - 1 |
| 182 | + full_data["map"][unit_no] = { |
| 183 | + "angle": [], |
| 184 | + "min_response": [], |
| 185 | + "max_response": [] |
| 186 | + } |
| 187 | + elif unit_no >= 0: |
| 188 | + result = re.match(r"([0-9]+), \[(.*), (.*)\]", line) |
| 189 | + if result is None: |
| 190 | + continue |
| 191 | + full_data["map"][unit_no]["angle"].append(float(result.group(1))) |
| 192 | + full_data["map"][unit_no]["min_response"].append(float(result.group(2))) |
| 193 | + full_data["map"][unit_no]["max_response"].append(float(result.group(3))) |
| 194 | + |
| 195 | + return full_data |
0 commit comments