import os import struct import bisect import numpy as np import _io import capnp from cereal import log as capnp_log class RawData(): def __init__(self, f): self.f = _io.FileIO(f, 'rb') self.lenn = struct.unpack("I", self.f.read(4))[0] self.count = os.path.getsize(f) / (self.lenn+4) def read(self, i): self.f.seek((self.lenn+4)*i + 4) return self.f.read(self.lenn) def yuv420_to_rgb(raw, image_dim=None, swizzled=False): def expand(x): x = np.repeat(x, 2, axis=0) return np.repeat(x, 2, axis=1) if image_dim is None: image_dim = (raw.shape[1]*2, raw.shape[2]*2) swizzled = True if not swizzled: img_data = np.array(raw, copy=False, dtype=np.uint8) uv_len = (image_dim[0]/2)*(image_dim[1]/2) img_data_u = expand(img_data[image_dim[0]*image_dim[1]: \ image_dim[0]*image_dim[1]+uv_len]. \ reshape(image_dim[0]/2, image_dim[1]/2)) img_data_v = expand(img_data[image_dim[0]*image_dim[1]+uv_len: \ image_dim[0]*image_dim[1]+2*uv_len]. \ reshape(image_dim[0]/2, image_dim[1]/2)) img_data_y = img_data[0:image_dim[0]*image_dim[1]].reshape(image_dim) else: img_data_y = np.zeros(image_dim, dtype=np.uint8) img_data_y[0::2, 0::2] = raw[0] img_data_y[1::2, 0::2] = raw[1] img_data_y[0::2, 1::2] = raw[2] img_data_y[1::2, 1::2] = raw[3] img_data_u = expand(raw[4]) img_data_v = expand(raw[5]) yuv = np.stack((img_data_y, img_data_u, img_data_v)).swapaxes(0,2).swapaxes(0,1) yuv = yuv.astype(np.int16) # http://maxsharabayko.blogspot.com/2016/01/fast-yuv-to-rgb-conversion-in-python-3.html # according to ITU-R BT.709 yuv[:,:, 0] = yuv[:,:, 0].clip(16, 235).astype(yuv.dtype) - 16 yuv[:,:,1:] = yuv[:,:,1:].clip(16, 240).astype(yuv.dtype) - 128 A = np.array([[1.164, 0.000, 1.793], [1.164, -0.213, -0.533], [1.164, 2.112, 0.000]]) # our result img = np.dot(yuv, A.T).clip(0, 255).astype('uint8') return img class YuvData(): def __init__(self, f, dim=(160,320)): self.f = _io.FileIO(f, 'rb') self.image_dim = dim self.image_size = self.image_dim[0]/2 * self.image_dim[1]/2 * 6 self.count = os.path.getsize(f) / self.image_size def read_frame(self, frame): self.f.seek(self.image_size*frame) raw = self.f.read(self.image_size) return raw def read_frames(self, range_start, range_len): self.f.seek(self.image_size*range_start) raw = self.f.read(self.image_size*range_len) return raw def read_frames_into(self, range_start, buf): self.f.seek(self.image_size*range_start) return self.f.readinto(buf) def read(self, frame): return yuv420_to_rgb(self.read_frame(frame), self.image_dim) def close(self): self.f.close() def __enter__(self): return self def __exit__(self, type, value, traceback): self.close() class OneReader(): def __init__(self, base_path, goofy=False, segment_range=None): self.base_path = base_path route_name = os.path.basename(base_path) self.rcamera_size = (304, 560) if segment_range is None: parent_path = os.path.dirname(base_path) self.segment_nums = [] for p in os.listdir(parent_path): if not p.startswith(route_name+"--"): continue self.segment_nums.append(int(p.rsplit("--", 1)[-1])) if not self.segment_nums: raise Exception("no route segments found") self.segment_nums.sort() self.segment_range = (self.segment_nums[0], self.segment_nums[-1]) else: self.segment_range = segment_range self.segment_nums = range(segment_range[0], segment_range[1]+1) for i in self.segment_nums: if not os.path.exists(base_path+"--"+str(i)): raise Exception("missing segment in provided range") # goofy data is broken with discontinuous logs if goofy and (self.segment_range[0] != 0 or self.segment_nums != range(self.segment_range[0], self.segment_range[1]+1)): raise Exception("goofy data needs all the segments for a route") self.cur_seg = None self.cur_seg_f = None # index the frames print("indexing frames {}...".format(self.segment_nums)) self.rcamera_encode_map = {} # frame_id -> (segment num, segment id, frame_time) last_frame_id = -1 if goofy: # goofy is goofy frame_size = self.rcamera_size[0]*self.rcamera_size[1]*3/2 # find the encode id ranges for each segment by using the rcamera file size segment_encode_ids = [] cur_encode_id = 0 for n in self.segment_nums: camera_path = os.path.join(self.seg_path(n), "rcamera") if not os.path.exists(camera_path): # for goofy, missing camera files means a bad route raise Exception("Missing camera file {}".format(camera_path)) camera_size = os.path.getsize(camera_path) assert (camera_size % frame_size) == 0 num_frames = camera_size / frame_size segment_encode_ids.append(cur_encode_id) cur_encode_id += num_frames last_encode_id = -1 # use the segment encode id map and frame events to build the frame index for n in self.segment_nums: log_path = os.path.join(self.seg_path(n), "rlog") if os.path.exists(log_path): with open(log_path, "rb") as f: for evt in capnp_log.Event.read_multiple(f): if evt.which() == 'frame': if evt.frame.frameId < last_frame_id: # a non-increasing frame id is bad route (eg visiond was restarted) raise Exception("non-increasing frame id") last_frame_id = evt.frame.frameId seg_i = bisect.bisect_right(segment_encode_ids, evt.frame.encodeId)-1 assert seg_i >= 0 seg_num = self.segment_nums[seg_i] seg_id = evt.frame.encodeId-segment_encode_ids[seg_i] frame_time = evt.logMonoTime / 1.0e9 self.rcamera_encode_map[evt.frame.frameId] = (seg_num, seg_id, frame_time) last_encode_id = evt.frame.encodeId if last_encode_id-cur_encode_id > 10: # too many missing frames is a bad route (eg route from before encoder rotating worked) raise Exception("goofy route is missing frames: {}, {}".format( last_encode_id, cur_encode_id)) else: # for harry data, build the index from encodeIdx events for n in self.segment_nums: log_path = os.path.join(self.seg_path(n), "rlog") if os.path.exists(log_path): with open(log_path, "rb") as f: for evt in capnp_log.Event.read_multiple(f): if evt.which() == 'encodeIdx' and evt.encodeIdx.type == 'bigBoxLossless': frame_time = evt.logMonoTime / 1.0e9 self.rcamera_encode_map[evt.encodeIdx.frameId] = ( evt.encodeIdx.segmentNum, evt.encodeIdx.segmentId, frame_time) print("done") # read the first event to find the start time self.reset_to_seg(self.segment_range[0]) for evt in self.events(): if evt.which() != 'initData': self.start_mono = evt.logMonoTime break self.reset_to_seg(self.segment_range[0]) def seg_path(self, num): return self.base_path+"--"+str(num) def reset_to_seg(self, seg): self.cur_seg = seg if self.cur_seg_f: self.cur_seg_f.close() self.cur_seg_f = None def seek_ts(self, ts): seek_seg = int(ts/60) if seek_seg < self.segment_range[0] or seek_seg > self.segment_range[1]: raise ValueError self.reset_to_seg(seek_seg) target_mono = self.start_mono + int(ts*1e9) for evt in self.events(): if evt.logMonoTime >= target_mono: break def read_event(self): while True: if self.cur_seg > self.segment_range[1]: return None if self.cur_seg_f is None: log_path = os.path.join(self.seg_path(self.cur_seg), "rlog") if not os.path.exists(log_path): print("missing log file!", log_path) self.cur_seg += 1 continue self.cur_seg_f = open(log_path, "rb") try: return capnp_log.Event.read(self.cur_seg_f) except capnp.lib.capnp.KjException as e: if 'EOF' in str(e): # dumb, but pycapnp does this too self.cur_seg_f.close() self.cur_seg_f = None self.cur_seg += 1 else: raise def events(self): while True: r = self.read_event() if r is None: break yield r def read_frame(self, frame_id): encode_idx = self.rcamera_encode_map.get(frame_id) if encode_idx is None: return None seg_num, seg_id, _ = encode_idx camera_path = os.path.join(self.seg_path(seg_num), "rcamera") if not os.path.exists(camera_path): return None with YuvData(camera_path, self.rcamera_size) as data: return data.read_frame(seg_id) def close(self): if self.cur_seg_f is not None: self.cur_seg_f.close() def __enter__(self): return self def __exit__(self, type, value, traceback): self.close()