extract_commaai_2k19.py

import numpy as np
import cv2
from extract import create_records, write_records
import glob
from multiprocessing import Pool
import camera
import orientation
import tensorflow as tf
import random

flags = tf.compat.v1.flags
FLAGS = flags.FLAGS

flags.DEFINE_integer("split", None, "Which split of the dataset to work on.")

def _bytes_feature(value):
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))

def _float_feature(value):
    return tf.train.Feature(float_list=tf.train.FloatList(value=[value]))

def _float_list_feature(values):
    return tf.train.Feature(float_list=tf.train.FloatList(value=values))

# takes in 2 ecef quaternions and returns the euler angle between them
def relative_orientation(ecef_quat_1, ecef_quat_2):
    ecef_mat_1 = orientation.quat2rot(ecef_quat_1)
    ecef_mat_2 = orientation.quat2rot(ecef_quat_2)

    relative_mat = np.matmul(ecef_mat_2, np.transpose(ecef_mat_1))

    return orientation.rot2euler(relative_mat)

# def find_lead_car(radar_packet):
#     lead_car = None
#     for ts, rr in radar_packet:
#         forward, left, relative = rr[:3]
#         if lead_car is None or abs(left) < abs(lead_car[2]):
#             lead_car = (ts, forward, left, relative)

#     return lead_car

# def get_lead_car(timestamps, radar_returns):
#     lead_cars = []
#     curr_group = [(timestamps[0], radar_returns[0])]

#     for ts, rr in zip(timestamps[1:], radar_returns[1:]):
#         if abs(ts - curr_group[-1][0]) < 0.02:
#             curr_group.append((ts, rr))
#         else:
#             lead_cars.append(find_lead_car(curr_group))
#             curr_group = [(ts, rr)]

#     return lead_cars

# def find_closest(lead_cars, frame_times):
#     closest_leads = []

#     for ft in frame_times:
#         closest = None
#         for lc in lead_cars:
#             if closest is None or abs(ft - lc[0]) < abs(ft - closest[0]):
#                 closest = lc
        
#         closest_leads.append(closest)

#     return closest_leads

def random_plus_or_minus():
    return np.random.rand() * 2 - 1

AUGMENTATION_STRENGTH = 0.25

def generate_pose_augment():
    forward = 1.0993402 * AUGMENTATION_STRENGTH * random_plus_or_minus()
    right = 0.03048608 * AUGMENTATION_STRENGTH * random_plus_or_minus()
    down = 0.04468765 * AUGMENTATION_STRENGTH * random_plus_or_minus()

    return np.array((forward, right, down))

def generate_orientation_augment():
    # a = 2e-2 * random_plus_or_minus()
    # b = 2e-2 * random_plus_or_minus()
    # c = 2e-2 * random_plus_or_minus()

    return np.array((0.0, 0.0, 0.0))

def augment_and_encode(frame, pose_augment, orientation_augment):
    frame = camera.transform_img(np.array(frame), augment_trans=camera.view_frame_from_device_frame.dot(pose_augment), augment_eulers=camera.view_frame_from_device_frame.dot(orientation_augment))

    frame = cv2.resize(frame, (640, 480), interpolation=cv2.INTER_AREA)
    ret, jpg = cv2.imencode(".jpg", frame)
    if not ret:
        raise Exception("couldn't encode the image")

    frame_bytes = jpg.tostring()

    return frame_bytes

def relative_pose_matrix(position_delta, angle_delta):
    angle_matrix = orientation.euler2rot(angle_delta)
    return np.vstack((np.hstack((angle_matrix, position_delta[:, None])), (0,0,0,1)))

def pose_components_from_matrix(pose_matrix):
    position_delta = pose_matrix[0:3, 3].T
    angle_delta = orientation.euler_from_rot(pose_matrix[0:3, 0:3])

    return position_delta, angle_delta

def extract_segment(path):
    frame_velocities =  np.linalg.norm(np.load(path + '/global_pose/frame_velocities'),axis=1)
    frame_velocities = list(map(lambda x: x.item(), frame_velocities))
    # frame_times = np.load(path + '/global_pose/frame_times')


    # radar_returns = np.load(path + '/processed_log/CAN/radar/value')
    # radar_time_stamps = np.load(path + '/processed_log/CAN/radar/t')

    # lead_cars = get_lead_car(radar_time_stamps, radar_returns)
    # matched_lead_cars = find_closest(lead_cars, frame_times)

    cam = cv2.VideoCapture(path + "/video.hevc")

    frame_positions = np.load(path + '/global_pose/frame_positions')
    frame_orientations = np.load(path + '/global_pose/frame_orientations')

    i = 0
    prev_frame = None
    examples = []
    while(True):
        ret, frame = cam.read()

        if ret:
            if prev_frame is not None:
                rel_position = camera.device_from_ecef(frame_positions[i-1], frame_orientations[i-1], frame_positions[i])
                rel_orientation = relative_orientation(frame_orientations[i-1], frame_orientations[i])
                rel_mat = relative_pose_matrix(rel_position, rel_orientation)

                pose_augment_one = generate_pose_augment()
                orientation_augment_one = generate_orientation_augment()
                mat_one = relative_pose_matrix(pose_augment_one, orientation_augment_one)

                frame_one_bytes = augment_and_encode(prev_frame, pose_augment_one, orientation_augment_one)

                pose_augment_two = generate_pose_augment()
                orientation_augment_two = generate_orientation_augment()
                mat_two = relative_pose_matrix(pose_augment_two, orientation_augment_one)

                frame_two_bytes = augment_and_encode(frame, pose_augment_two, orientation_augment_two)

                total_mat = rel_mat.dot(np.linalg.inv(mat_one)).dot(mat_two)
                total_position, total_orientation = pose_components_from_matrix(total_mat)

                example = tf.train.Example(features=tf.train.Features(feature={
                    'frame_one': _bytes_feature(frame_one_bytes),
                    'frame_two': _bytes_feature(frame_two_bytes),
                    'plus_one_position': _float_list_feature(total_position),
                    'plus_one_orientation': _float_list_feature(total_orientation),
                    # 'lead_car_one': _float_list_feature(matched_lead_cars[i][1:]),
                    # # minus relative delta, because backward in time
                    # 'lead_car_two': _float_list_feature((1, 1, -1) * matched_lead_cars[i+1][1:]),
                    'speed': _float_feature(20 * np.linalg.norm(total_position)),
                }))
                examples.append(example.SerializeToString())


            prev_frame = frame
            i += 1
        else:
            break

    return examples

def write_tfrecord(dirs, out):
    with tf.io.TFRecordWriter(out) as writer:
        example_buffer = []
        for i, dir in enumerate(dirs):
            example_buffer.extend(extract_segment(dir))
            print("extracted {} segments".format(i))

        random.shuffle(example_buffer)
        for e in example_buffer:
            writer.write(e)


def write_to_file(dirs, file):
    with open(file, "w") as f:
        f.write("\n".join(dirs))

# segments = glob.glob("/mnt/e/commaai/comma2k19/*/*/*")
# random.shuffle(segments)

# train = segments[:1933]
# validation = segments[1933:]

# write_to_file(train, "2k19_train_split.txt")
# write_to_file(validation, "2k19_val_split.txt")

# write_tfrecord(["/mnt/d/commaai/comma2k19/Chunk_1/b0c9d2329ad1606b_2018-07-27--06-03-57/4"], "test.tfrecord")

dirs = None
with open("2k19_train_split.txt") as f:
    dirs = [x.replace("/mnt/e", "/mnt/d") for x in f.read().splitlines()]

write_tfrecord(dirs[FLAGS.split::100], "/mnt/f/Training Data/commaai/2k19_train_augmented/{}.tfrecord".format(FLAGS.split))