extract_calib.py

import cv2
import tensorflow as tf
import numpy as np

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))

def write_records(examples, path):
    writer = tf.python_io.TFRecordWriter(path)
    for e in examples:
        writer.write(e)

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

    image_bytes = jpg.tostring()
    example = tf.train.Example(features=tf.train.Features(feature={
        'image_raw': _bytes_feature(image_bytes),
        'label': _float_feature(speed)
    }))

    return example.SerializeToString()

def create_records(cam, speeds):
    current_frame = 0
    examples = []

    while(True):
        ret, frame = cam.read()

        if ret:
            speed = speeds[current_frame]

            examples.append(create_record(frame, speed))

            current_frame += 1
        else:
            break

    return examples

DIRECTORY = "/home/lucas/Development/calib_challenge/labeled"

def main():
    video = "4"
    cam = cv2.VideoCapture("{}/{}.hevc".format(DIRECTORY, video))
    labels = open("{}/{}.txt".format(DIRECTORY, video), 'r').readlines()

    frames = []

    while(True):
        ret, frame = cam.read()

        if ret:
            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()
            frames.append(frame_bytes)

        else:
            break

    pitches = []
    yaws = []

    for l in labels:
        pitch, yaw = l.split()
        pitches.append(float(pitch))
        yaws.append(float(yaw))

    assert len(frames) == len(pitches) == len(yaws)
    examples = []
    for i in range(len(frames[:-1])):
        example = tf.train.Example(features=tf.train.Features(feature={
            'frame_one': _bytes_feature(frames[i]),
            'frame_two': _bytes_feature(frames[i+1]),
            'plus_one_position': _float_list_feature([0.0, 0.0, 0.0]),
            'plus_one_orientation': _float_list_feature([0.0, 0.0, 0.0]),
            'pitch': _float_feature(pitches[i]),
            'yaw': _float_feature(yaws[i]),
            'speed': _float_feature(0.0),
        }))
        examples.append(example.SerializeToString())


    with tf.io.TFRecordWriter("/mnt/Datasets/calib_{}.tfrecord".format(video)) as writer:
        for e in examples:
            writer.write(e)



if __name__ == '__main__':
    main()