From 99c5800ea5bb0704d96a6e459af2acf1a1050e3c Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 17 Jan 2020 10:07:22 -0800 Subject: [PATCH] merge in tools old-commit-hash: 29ac3da7b84426c6764150cb233b9c7bab1446d8 --- tools/LICENSE | 21 + tools/README.md | 294 ++++++ tools/__init__.py | 0 tools/carcontrols/__init__.py | 0 tools/carcontrols/debug_controls.py | 100 +++ tools/carcontrols/joystick_test.py | 125 +++ tools/carcontrols/joystickd.py | 68 ++ tools/clib/.gitignore | 1 + tools/clib/FrameReader.cpp | 176 ++++ tools/clib/FrameReader.hpp | 58 ++ tools/clib/SConscript | 8 + tools/clib/cframereader.pyx | 20 + tools/clib/channel.hpp | 35 + tools/lib/__init__.py | 0 tools/lib/async_generator.py | 352 ++++++++ tools/lib/cache.py | 9 + tools/lib/exceptions.py | 2 + tools/lib/file_helpers.py | 23 + tools/lib/filereader.py | 3 + tools/lib/framereader.py | 1284 +++++++++++++++++++++++++++ tools/lib/index_log/.gitignore | 1 + tools/lib/index_log/Makefile | 19 + tools/lib/index_log/index_log.cc | 66 ++ tools/lib/kbhit.py | 92 ++ tools/lib/lazy_property.py | 12 + tools/lib/log_util.py | 111 +++ tools/lib/logreader.py | 205 +++++ tools/lib/mkvparse/README.md | 24 + tools/lib/mkvparse/__init__.py | 0 tools/lib/mkvparse/mkvgen.py | 188 ++++ tools/lib/mkvparse/mkvindex.py | 87 ++ tools/lib/mkvparse/mkvparse.py | 761 ++++++++++++++++ tools/lib/pollable_queue.py | 107 +++ tools/lib/route.py | 97 ++ tools/lib/route_framereader.py | 86 ++ tools/lib/tests/test_readers.py | 56 ++ tools/lib/vidindex/.gitignore | 1 + tools/lib/vidindex/Makefile | 6 + tools/lib/vidindex/bitstream.c | 118 +++ tools/lib/vidindex/bitstream.h | 26 + tools/lib/vidindex/vidindex.c | 307 +++++++ tools/livedm/helpers.py | 30 + tools/livedm/livedm.py | 79 ++ tools/misc/save_ubloxraw_stream.py | 59 ++ tools/nui/.gitignore | 8 + tools/nui/FileReader.cpp | 138 +++ tools/nui/FileReader.hpp | 68 ++ tools/nui/README | 9 + tools/nui/Unlogger.cpp | 182 ++++ tools/nui/Unlogger.hpp | 36 + tools/nui/build.sh | 4 + tools/nui/main.cpp | 216 +++++ tools/nui/nui.pro | 3 + tools/nui/test/.gitignore | 1 + tools/nui/test/TestFrameReader.cpp | 14 + tools/nui/test/TestFrameReader.hpp | 8 + tools/nui/test/test.pro | 3 + tools/replay/__init__.py | 0 tools/replay/camera.py | 112 +++ tools/replay/lib/__init__.py | 0 tools/replay/lib/ui_helpers.py | 314 +++++++ tools/replay/mapd.py | 68 ++ tools/replay/rqplot.py | 83 ++ tools/replay/ui.py | 278 ++++++ tools/replay/unlogger.py | 446 ++++++++++ tools/requirements.txt | 11 + tools/sim/.gitignore | 3 + tools/sim/can.py | 100 +++ tools/sim/controller.py | 155 ++++ tools/sim/get.sh | 10 + tools/sim/replay.sh | 5 + tools/sim/start.sh | 4 + tools/ssh/config | 9 + tools/ssh/key/id_rsa | 28 + tools/ssh/via-smays.sh | 3 + tools/ssh/via-wifi.sh | 1 + tools/steer.gif | 3 + tools/stream.gif | 3 + tools/streamer/streamerd.py | 93 ++ 79 files changed, 7536 insertions(+) create mode 100644 tools/LICENSE create mode 100644 tools/README.md create mode 100644 tools/__init__.py create mode 100644 tools/carcontrols/__init__.py create mode 100755 tools/carcontrols/debug_controls.py create mode 100755 tools/carcontrols/joystick_test.py create mode 100755 tools/carcontrols/joystickd.py create mode 100644 tools/clib/.gitignore create mode 100644 tools/clib/FrameReader.cpp create mode 100644 tools/clib/FrameReader.hpp create mode 100644 tools/clib/SConscript create mode 100644 tools/clib/cframereader.pyx create mode 100644 tools/clib/channel.hpp create mode 100644 tools/lib/__init__.py create mode 100644 tools/lib/async_generator.py create mode 100644 tools/lib/cache.py create mode 100644 tools/lib/exceptions.py create mode 100644 tools/lib/file_helpers.py create mode 100644 tools/lib/filereader.py create mode 100644 tools/lib/framereader.py create mode 100644 tools/lib/index_log/.gitignore create mode 100644 tools/lib/index_log/Makefile create mode 100644 tools/lib/index_log/index_log.cc create mode 100644 tools/lib/kbhit.py create mode 100644 tools/lib/lazy_property.py create mode 100755 tools/lib/log_util.py create mode 100644 tools/lib/logreader.py create mode 100644 tools/lib/mkvparse/README.md create mode 100644 tools/lib/mkvparse/__init__.py create mode 100755 tools/lib/mkvparse/mkvgen.py create mode 100644 tools/lib/mkvparse/mkvindex.py create mode 100644 tools/lib/mkvparse/mkvparse.py create mode 100644 tools/lib/pollable_queue.py create mode 100644 tools/lib/route.py create mode 100644 tools/lib/route_framereader.py create mode 100755 tools/lib/tests/test_readers.py create mode 100644 tools/lib/vidindex/.gitignore create mode 100644 tools/lib/vidindex/Makefile create mode 100644 tools/lib/vidindex/bitstream.c create mode 100644 tools/lib/vidindex/bitstream.h create mode 100644 tools/lib/vidindex/vidindex.c create mode 100644 tools/livedm/helpers.py create mode 100644 tools/livedm/livedm.py create mode 100644 tools/misc/save_ubloxraw_stream.py create mode 100644 tools/nui/.gitignore create mode 100644 tools/nui/FileReader.cpp create mode 100644 tools/nui/FileReader.hpp create mode 100644 tools/nui/README create mode 100644 tools/nui/Unlogger.cpp create mode 100644 tools/nui/Unlogger.hpp create mode 100755 tools/nui/build.sh create mode 100644 tools/nui/main.cpp create mode 100644 tools/nui/nui.pro create mode 100644 tools/nui/test/.gitignore create mode 100644 tools/nui/test/TestFrameReader.cpp create mode 100644 tools/nui/test/TestFrameReader.hpp create mode 100644 tools/nui/test/test.pro create mode 100644 tools/replay/__init__.py create mode 100755 tools/replay/camera.py create mode 100644 tools/replay/lib/__init__.py create mode 100644 tools/replay/lib/ui_helpers.py create mode 100755 tools/replay/mapd.py create mode 100755 tools/replay/rqplot.py create mode 100755 tools/replay/ui.py create mode 100755 tools/replay/unlogger.py create mode 100644 tools/requirements.txt create mode 100644 tools/sim/.gitignore create mode 100755 tools/sim/can.py create mode 100755 tools/sim/controller.py create mode 100755 tools/sim/get.sh create mode 100755 tools/sim/replay.sh create mode 100755 tools/sim/start.sh create mode 100644 tools/ssh/config create mode 100644 tools/ssh/key/id_rsa create mode 100755 tools/ssh/via-smays.sh create mode 100755 tools/ssh/via-wifi.sh create mode 100644 tools/steer.gif create mode 100644 tools/stream.gif create mode 100755 tools/streamer/streamerd.py diff --git a/tools/LICENSE b/tools/LICENSE new file mode 100644 index 0000000000..b8fd9e299c --- /dev/null +++ b/tools/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2018 comma.ai + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/tools/README.md b/tools/README.md new file mode 100644 index 0000000000..e5ffb15c9c --- /dev/null +++ b/tools/README.md @@ -0,0 +1,294 @@ +openpilot-tools +============ + +Repo which contains tools to facilitate development and debugging of [openpilot](openpilot.comma.ai). + +![Imgur](https://i.imgur.com/IdfBgwK.jpg) + + +Table of Contents +============ + + + * [Requirements](#requirements) + * [Setup](#setup) + * [Tool examples](#tool-examples) + * [Replay driving data](#replay-driving-data) + * [Debug car controls](#debug-car-controls) + * [Stream replayed CAN messages to EON](#stream-replayed-can-messages-to-eon) + * [Stream EON video data to a PC](#stream-eon-video-data-to-a-pc) + * [Welcomed contributions](#welcomed-contributions) + + + +Requirements +============ + +openpilot-tools and the following setup steps are developed and tested on Ubuntu 16.04, MacOS 10.14.2 and Python 3.7.3. + +Setup +============ + + +1. Install native dependencies (Mac and Ubuntu sections listed below) + + **Ubuntu** + + - core tools + ```bash + sudo apt install git curl python-pip + sudo pip install --upgrade pip>=18.0 pipenv + ``` + + - ffmpeg (tested with 3.3.2) + ```bash + sudo apt install ffmpeg libavformat-dev libavcodec-dev libavdevice-dev libavutil-dev libswscale-dev libavresample-dev libavfilter-dev + ``` + + - build tools + ```bash + sudo apt install autoconf automake clang clang-3.8 libtool pkg-config build-essential + ``` + + - libarchive-dev (tested with 3.1.2-11ubuntu0.16.04.4) + ```bash + sudo apt install libarchive-dev + ``` + + - qt python binding (tested with python-qt4, 4.11.4+dfsg-1build4) + ```bash + sudo apt install python-qt4 + ``` + + - zmq 4.2.3 (required for replay) + ```bash + curl -LO https://github.com/zeromq/libzmq/releases/download/v4.2.3/zeromq-4.2.3.tar.gz + tar xfz zeromq-4.2.3.tar.gz + cd zeromq-4.2.3 + ./autogen.sh + ./configure CPPFLAGS=-DPIC CFLAGS=-fPIC CXXFLAGS=-fPIC LDFLAGS=-fPIC --disable-shared --enable-static + make + sudo make install + ``` + + **Mac** + + - brew + ``` bash + /usr/bin/ruby -e "$(curl -fsSL https://raw.githubusercontent.com/Homebrew/install/master/install)" + ``` + + - core tools + ``` bash + brew install git + sudo pip install --upgrade pip pipenv + xcode-select --install + ``` + + - ffmpeg (tested with 3.4.1) + ```bash + brew install ffmpeg + ``` + + - build tools + ```bash + brew install autoconf automake libtool llvm pkg-config + ``` + + - libarchive-dev (tested with 3.3.3) + ```bash + brew install libarchive + ``` + + - qt for Mac + ```bash + brew install qt + ``` + + - zmq 4.3.1 (required for replay) + ```bash + brew install zeromq + ``` + +2. Install Cap'n Proto + + ```bash + curl -O https://capnproto.org/capnproto-c++-0.6.1.tar.gz + tar xvf capnproto-c++-0.6.1.tar.gz + cd capnproto-c++-0.6.1 + ./configure --prefix=/usr/local CPPFLAGS=-DPIC CFLAGS=-fPIC CXXFLAGS=-fPIC LDFLAGS=-fPIC --disable-shared --enable-static + make -j4 + sudo make install + + cd .. + git clone https://github.com/commaai/c-capnproto.git + cd c-capnproto + git submodule update --init --recursive + autoreconf -f -i -s + CFLAGS="-fPIC" ./configure --prefix=/usr/local + make -j4 + sudo make install + ``` + + + + +2. Clone openpilot if you haven't already + + ```bash + git clone https://github.com/commaai/openpilot.git + cd openpilot + pipenv install # Install dependencies in a virtualenv + pipenv shell # Activate the virtualenv + ``` + + **For Mac users** + + Recompile longitudinal_mpc for mac + + Navigate to: + ``` bash + cd selfdrive/controls/lib/longitudinal_mpc + make clean + make + ``` + +3. Clone tools within openpilot, and install dependencies + + ```bash + git clone https://github.com/commaai/openpilot-tools.git tools + cd tools + git checkout # the tag must match the openpilot version you are using (see https://github.com/commaai/openpilot-tools/tags) + pip install -r requirements.txt # Install openpilot-tools dependencies in virtualenv + ``` + +4. Add openpilot to your `PYTHONPATH`. + + For bash users: + ```bash + echo 'export PYTHONPATH="$PYTHONPATH:"' >> ~/.bashrc + source ~/.bashrc + ``` + +5. Add some folders to root + ```bash + sudo mkdir /data + sudo mkdir /data/params + sudo chown $USER /data/params + ``` + +6. Try out some tools! + + +Tool examples +============ + + +Replay driving data +------------- + +**Hardware needed**: none + +`unlogger.py` replays data collected with [chffrplus](https://github.com/commaai/chffrplus) or [openpilot](https://github.com/commaai/openpilot). + +You'll need to download log and camera files into a local directory. Download these from the footer of the comma [explorer](https://my.comma.ai) or SCP from your device. + +Usage: + +``` +python replay/unlogger.py + +#Example: + +#python replay/unlogger.py '99c94dc769b5d96e|2018-11-14--13-31-42' /home/batman/unlogger_data + +#Within /home/batman/unlogger_data: +# 99c94dc769b5d96e|2018-11-14--13-31-42--0--fcamera.hevc +# 99c94dc769b5d96e|2018-11-14--13-31-42--0--rlog.bz2 +# ... + +# In another terminal you can run a debug visualizer: +python replay/ui.py # Define the environmental variable HORIZONTAL is the ui layout is too tall +``` +![Imgur](https://i.imgur.com/Yppe0h2.png) + + +Debug car controls +------------- + +**Hardware needed**: [panda](panda.comma.ai), [giraffe](https://comma.ai/shop/products/giraffe/), joystick + +Use the panda's OBD-II port to connect with your car and a usb cable to connect the panda to your pc. +Also, connect a joystick to your pc. + +`joystickd.py` runs a deamon that reads inputs from a joystick and publishes them over zmq. +`boardd.py` sends the CAN messages from your pc to the panda. +`debug_controls` is a mocked version of `controlsd.py` and uses input from a joystick to send controls to your car. + +Usage: +``` +python carcontrols/joystickd.py + +# In another terminal: +selfdrive/boardd/tests/boardd_old.py # Make sure the safety setting is hardcoded to ALL_OUTPUT + +# In another terminal: +python carcontrols/debug_controls.py + +``` +![Imgur](steer.gif) + + +Stream replayed CAN messages to EON +------------- + +**Hardware needed**: 2 x [panda](panda.comma.ai), [debug board](https://comma.ai/shop/products/panda-debug-board/), [EON](https://comma.ai/shop/products/eon-gold-dashcam-devkit/). + +It is possible to replay CAN messages as they were recorded and forward them to a EON.  +Connect 2 pandas to the debug board. A panda connects to the PC, the other panda connects to the EON. + +Usage: +``` +# With MOCK=1 boardd will read logged can messages from a replay and send them to the panda. +MOCK=1 selfdrive/boardd/tests/boardd_old.py + +# In another terminal: +python replay/unlogger.py + +``` +![Imgur](https://i.imgur.com/AcurZk8.jpg) + + +Stream EON video data to a PC +------------- + +**Hardware needed**: [EON](https://comma.ai/shop/products/eon-gold-dashcam-devkit/), [comma Smays](https://comma.ai/shop/products/comma-smays-adapter/). + +You can connect your EON to your pc using the Ethernet cable provided with the comma Smays and you'll be able to stream data from your EON, in real time, with low latency. A useful application is being able to stream the raw video frames at 20fps, as captured by the EON's camera. + +Usage: +``` +# ssh into the eon and run loggerd with the flag "--stream". In ../selfdrive/manager.py you can change: +# ... +# "loggerd": ("selfdrive/loggerd", ["./loggerd"]), +# ... +# with: +# ... +# "loggerd": ("selfdrive/loggerd", ["./loggerd", "--stream"]), +# ... + +# On the PC: +# To receive frames from the EON and re-publish them. Set PYGAME env variable if you want to display the video stream +python streamer/streamerd.py +``` + +![Imgur](stream.gif) + + +Welcomed contributions +============= + +* Documentation: code comments, better tutorials, etc.. +* Support for other platforms other than Ubuntu 16.04. +* Performance improvements: the tools have been developed on high-performance workstations (12+ logical cores with 32+ GB of RAM), so they are not optimized for running efficiently. For example, `ui.py` might not be able to run real-time on most PCs. +* More tools: anything that you think might be helpful to others. diff --git a/tools/__init__.py b/tools/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tools/carcontrols/__init__.py b/tools/carcontrols/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tools/carcontrols/debug_controls.py b/tools/carcontrols/debug_controls.py new file mode 100755 index 0000000000..527ef42fc0 --- /dev/null +++ b/tools/carcontrols/debug_controls.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python +import struct +from common.numpy_fast import clip +from common.params import Params +from copy import copy +from cereal import car, log +import cereal.messaging as messaging +from selfdrive.car.car_helpers import get_car +from selfdrive.boardd.boardd import can_list_to_can_capnp + +HwType = log.HealthData.HwType + + +def steer_thread(): + poller = messaging.Poller() + + logcan = messaging.sub_sock('can') + health = messaging.sub_sock('health') + joystick_sock = messaging.sub_sock('testJoystick', conflate=True, poller=poller) + + carstate = messaging.pub_sock('carState') + carcontrol = messaging.pub_sock('carControl') + sendcan = messaging.pub_sock('sendcan') + + button_1_last = 0 + enabled = False + + # wait for health and CAN packets + hw_type = messaging.recv_one(health).health.hwType + has_relay = hw_type in [HwType.blackPanda, HwType.uno] + print("Waiting for CAN messages...") + messaging.get_one_can(logcan) + + CI, CP = get_car(logcan, sendcan, has_relay) + Params().put("CarParams", CP.to_bytes()) + + CC = car.CarControl.new_message() + + while True: + + # send + joystick = messaging.recv_one(joystick_sock) + can_strs = messaging.drain_sock_raw(logcan, wait_for_one=True) + CS = CI.update(CC, can_strs) + + # Usually axis run in pairs, up/down for one, and left/right for + # the other. + actuators = car.CarControl.Actuators.new_message() + + if joystick is not None: + axis_3 = clip(-joystick.testJoystick.axes[3] * 1.05, -1., 1.) # -1 to 1 + actuators.steer = axis_3 + actuators.steerAngle = axis_3 * 43. # deg + axis_1 = clip(-joystick.testJoystick.axes[1] * 1.05, -1., 1.) # -1 to 1 + actuators.gas = max(axis_1, 0.) + actuators.brake = max(-axis_1, 0.) + + pcm_cancel_cmd = joystick.testJoystick.buttons[0] + button_1 = joystick.testJoystick.buttons[1] + if button_1 and not button_1_last: + enabled = not enabled + + button_1_last = button_1 + + #print "enable", enabled, "steer", actuators.steer, "accel", actuators.gas - actuators.brake + + hud_alert = 0 + audible_alert = 0 + if joystick.testJoystick.buttons[2]: + audible_alert = "beepSingle" + if joystick.testJoystick.buttons[3]: + audible_alert = "chimeRepeated" + hud_alert = "steerRequired" + + CC.actuators.gas = actuators.gas + CC.actuators.brake = actuators.brake + CC.actuators.steer = actuators.steer + CC.actuators.steerAngle = actuators.steerAngle + CC.hudControl.visualAlert = hud_alert + CC.hudControl.setSpeed = 20 + CC.cruiseControl.cancel = pcm_cancel_cmd + CC.enabled = enabled + can_sends = CI.apply(CC) + sendcan.send(can_list_to_can_capnp(can_sends, msgtype='sendcan')) + + # broadcast carState + cs_send = messaging.new_message() + cs_send.init('carState') + cs_send.carState = copy(CS) + carstate.send(cs_send.to_bytes()) + + # broadcast carControl + cc_send = messaging.new_message() + cc_send.init('carControl') + cc_send.carControl = copy(CC) + carcontrol.send(cc_send.to_bytes()) + + +if __name__ == "__main__": + steer_thread() diff --git a/tools/carcontrols/joystick_test.py b/tools/carcontrols/joystick_test.py new file mode 100755 index 0000000000..e031fb506f --- /dev/null +++ b/tools/carcontrols/joystick_test.py @@ -0,0 +1,125 @@ +#!/usr/bin/env python +import pygame + +# Define some colors +BLACK = ( 0, 0, 0) +WHITE = ( 255, 255, 255) + +# This is a simple class that will help us print to the screen +# It has nothing to do with the joysticks, just outputting the +# information. +class TextPrint: + def __init__(self): + self.reset() + self.font = pygame.font.Font(None, 20) + + def printf(self, screen, textString): + textBitmap = self.font.render(textString, True, BLACK) + screen.blit(textBitmap, [self.x, self.y]) + self.y += self.line_height + + def reset(self): + self.x = 10 + self.y = 10 + self.line_height = 15 + + def indent(self): + self.x += 10 + + def unindent(self): + self.x -= 10 + + +pygame.init() + +# Set the width and height of the screen [width,height] +size = [500, 700] +screen = pygame.display.set_mode(size) + +pygame.display.set_caption("My Game") + +#Loop until the user clicks the close button. +done = False + +# Used to manage how fast the screen updates +clock = pygame.time.Clock() + +# Initialize the joysticks +pygame.joystick.init() + +# Get ready to print +textPrint = TextPrint() + +# -------- Main Program Loop ----------- +while done==False: + # EVENT PROCESSING STEP + for event in pygame.event.get(): # User did something + if event.type == pygame.QUIT: # If user clicked close + done=True # Flag that we are done so we exit this loop + + # Possible joystick actions: JOYAXISMOTION JOYBALLMOTION JOYBUTTONDOWN JOYBUTTONUP JOYHATMOTION + if event.type == pygame.JOYBUTTONDOWN: + print("Joystick button pressed.") + if event.type == pygame.JOYBUTTONUP: + print("Joystick button released.") + + + # DRAWING STEP + # First, clear the screen to white. Don't put other drawing commands + # above this, or they will be erased with this command. + screen.fill(WHITE) + textPrint.reset() + + # Get count of joysticks + joystick_count = pygame.joystick.get_count() + + textPrint.printf(screen, "Number of joysticks: {}".format(joystick_count) ) + textPrint.indent() + + # For each joystick: + joystick = pygame.joystick.Joystick(0) + joystick.init() + + textPrint.printf(screen, "Joystick {}".format(0) ) + textPrint.indent() + + # Get the name from the OS for the controller/joystick + name = joystick.get_name() + textPrint.printf(screen, "Joystick name: {}".format(name) ) + + # Usually axis run in pairs, up/down for one, and left/right for + # the other. + axes = joystick.get_numaxes() + textPrint.printf(screen, "Number of axes: {}".format(axes) ) + textPrint.indent() + + for i in range( axes ): + axis = joystick.get_axis( i ) + textPrint.printf(screen, "Axis {} value: {:>6.3f}".format(i, axis) ) + textPrint.unindent() + + buttons = joystick.get_numbuttons() + textPrint.printf(screen, "Number of buttons: {}".format(buttons) ) + textPrint.indent() + + for i in range( buttons ): + button = joystick.get_button( i ) + textPrint.printf(screen, "Button {:>2} value: {}".format(i,button) ) + textPrint.unindent() + + + textPrint.unindent() + + + # ALL CODE TO DRAW SHOULD GO ABOVE THIS COMMENT + + # Go ahead and update the screen with what we've drawn. + pygame.display.flip() + + # Limit to 20 frames per second + clock.tick(20) + +# Close the window and quit. +# If you forget this line, the program will 'hang' +# on exit if running from IDLE. +pygame.quit () diff --git a/tools/carcontrols/joystickd.py b/tools/carcontrols/joystickd.py new file mode 100755 index 0000000000..c80deec606 --- /dev/null +++ b/tools/carcontrols/joystickd.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python + +# This process publishes joystick events. Such events can be suscribed by +# mocked car controller scripts. + + +### this process needs pygame and can't run on the EON ### + +import pygame +import zmq +import cereal.messaging as messaging + + +def joystick_thread(): + joystick_sock = messaging.pub_sock('testJoystick') + + pygame.init() + + # Used to manage how fast the screen updates + clock = pygame.time.Clock() + + # Initialize the joysticks + pygame.joystick.init() + + # Get count of joysticks + joystick_count = pygame.joystick.get_count() + if joystick_count > 1: + raise ValueError("More than one joystick attached") + elif joystick_count < 1: + raise ValueError("No joystick found") + + # -------- Main Program Loop ----------- + while True: + # EVENT PROCESSING STEP + for event in pygame.event.get(): # User did something + if event.type == pygame.QUIT: # If user clicked close + pass + # Available joystick events: JOYAXISMOTION JOYBALLMOTION JOYBUTTONDOWN JOYBUTTONUP JOYHATMOTION + if event.type == pygame.JOYBUTTONDOWN: + print("Joystick button pressed.") + if event.type == pygame.JOYBUTTONUP: + print("Joystick button released.") + + joystick = pygame.joystick.Joystick(0) + joystick.init() + + # Usually axis run in pairs, up/down for one, and left/right for + # the other. + axes = [] + buttons = [] + + for a in range(joystick.get_numaxes()): + axes.append(joystick.get_axis(a)) + + for b in range(joystick.get_numbuttons()): + buttons.append(bool(joystick.get_button(b))) + + dat = messaging.new_message() + dat.init('testJoystick') + dat.testJoystick.axes = axes + dat.testJoystick.buttons = buttons + joystick_sock.send(dat.to_bytes()) + + # Limit to 100 frames per second + clock.tick(100) + +if __name__ == "__main__": + joystick_thread() diff --git a/tools/clib/.gitignore b/tools/clib/.gitignore new file mode 100644 index 0000000000..ae8a795a92 --- /dev/null +++ b/tools/clib/.gitignore @@ -0,0 +1 @@ +cframereader.cpp diff --git a/tools/clib/FrameReader.cpp b/tools/clib/FrameReader.cpp new file mode 100644 index 0000000000..91d31deec2 --- /dev/null +++ b/tools/clib/FrameReader.cpp @@ -0,0 +1,176 @@ +#include "FrameReader.hpp" +#include +#include + +static int ffmpeg_lockmgr_cb(void **arg, enum AVLockOp op) { + pthread_mutex_t *mutex = (pthread_mutex_t *)*arg; + int err; + + switch (op) { + case AV_LOCK_CREATE: + mutex = (pthread_mutex_t *)malloc(sizeof(*mutex)); + if (!mutex) + return AVERROR(ENOMEM); + if ((err = pthread_mutex_init(mutex, NULL))) { + free(mutex); + return AVERROR(err); + } + *arg = mutex; + return 0; + case AV_LOCK_OBTAIN: + if ((err = pthread_mutex_lock(mutex))) + return AVERROR(err); + + return 0; + case AV_LOCK_RELEASE: + if ((err = pthread_mutex_unlock(mutex))) + return AVERROR(err); + + return 0; + case AV_LOCK_DESTROY: + if (mutex) + pthread_mutex_destroy(mutex); + free(mutex); + *arg = NULL; + return 0; + } + return 1; +} + +FrameReader::FrameReader(const char *fn) { + int ret; + + ret = av_lockmgr_register(ffmpeg_lockmgr_cb); + assert(ret >= 0); + + avformat_network_init(); + av_register_all(); + + snprintf(url, sizeof(url)-1, "http://data.comma.life/%s", fn); + t = new std::thread([&]() { this->loaderThread(); }); +} + +void FrameReader::loaderThread() { + int ret; + + if (avformat_open_input(&pFormatCtx, url, NULL, NULL) != 0) { + fprintf(stderr, "error loading %s\n", url); + valid = false; + return; + } + av_dump_format(pFormatCtx, 0, url, 0); + + auto pCodecCtxOrig = pFormatCtx->streams[0]->codec; + auto pCodec = avcodec_find_decoder(pCodecCtxOrig->codec_id); + assert(pCodec != NULL); + + pCodecCtx = avcodec_alloc_context3(pCodec); + ret = avcodec_copy_context(pCodecCtx, pCodecCtxOrig); + assert(ret == 0); + + ret = avcodec_open2(pCodecCtx, pCodec, NULL); + assert(ret >= 0); + + sws_ctx = sws_getContext(width, height, AV_PIX_FMT_YUV420P, + width, height, AV_PIX_FMT_BGR24, + SWS_BILINEAR, NULL, NULL, NULL); + assert(sws_ctx != NULL); + + AVPacket *pkt = (AVPacket *)malloc(sizeof(AVPacket)); + assert(pkt != NULL); + bool first = true; + while (av_read_frame(pFormatCtx, pkt)>=0) { + //printf("%d pkt %d %d\n", pkts.size(), pkt->size, pkt->pos); + if (first) { + AVFrame *pFrame = av_frame_alloc(); + int frameFinished; + avcodec_decode_video2(pCodecCtx, pFrame, &frameFinished, pkt); + first = false; + } + pkts.push_back(pkt); + pkt = (AVPacket *)malloc(sizeof(AVPacket)); + assert(pkt != NULL); + } + free(pkt); + printf("framereader download done\n"); + joined = true; + + // cache + while (1) { + GOPCache(to_cache.get()); + } +} + + +void FrameReader::GOPCache(int idx) { + AVFrame *pFrame; + int gop = idx - idx%15; + + mcache.lock(); + bool has_gop = cache.find(gop) != cache.end(); + mcache.unlock(); + + if (!has_gop) { + //printf("caching %d\n", gop); + for (int i = gop; i < gop+15; i++) { + if (i >= pkts.size()) break; + //printf("decode %d\n", i); + int frameFinished; + pFrame = av_frame_alloc(); + avcodec_decode_video2(pCodecCtx, pFrame, &frameFinished, pkts[i]); + uint8_t *dat = toRGB(pFrame)->data[0]; + mcache.lock(); + cache.insert(std::make_pair(i, dat)); + mcache.unlock(); + } + } +} + +AVFrame *FrameReader::toRGB(AVFrame *pFrame) { + AVFrame *pFrameRGB = av_frame_alloc(); + int numBytes = avpicture_get_size(AV_PIX_FMT_BGR24, pFrame->width, pFrame->height); + uint8_t *buffer = (uint8_t *)av_malloc(numBytes*sizeof(uint8_t)); + avpicture_fill((AVPicture *)pFrameRGB, buffer, AV_PIX_FMT_BGR24, pFrame->width, pFrame->height); + sws_scale(sws_ctx, (uint8_t const * const *)pFrame->data, + pFrame->linesize, 0, pFrame->height, + pFrameRGB->data, pFrameRGB->linesize); + return pFrameRGB; +} + +uint8_t *FrameReader::get(int idx) { + if (!valid) return NULL; + waitForReady(); + // TODO: one line? + uint8_t *dat = NULL; + + // lookahead + to_cache.put(idx); + to_cache.put(idx+15); + + mcache.lock(); + auto it = cache.find(idx); + if (it != cache.end()) { + dat = it->second; + } + mcache.unlock(); + + if (dat == NULL) { + to_cache.put_front(idx); + // lookahead + while (dat == NULL) { + // wait for frame + usleep(50*1000); + // check for frame + mcache.lock(); + auto it = cache.find(idx); + if (it != cache.end()) dat = it->second; + mcache.unlock(); + if (dat == NULL) { + printf("."); + fflush(stdout); + } + } + } + return dat; +} + diff --git a/tools/clib/FrameReader.hpp b/tools/clib/FrameReader.hpp new file mode 100644 index 0000000000..db655d641d --- /dev/null +++ b/tools/clib/FrameReader.hpp @@ -0,0 +1,58 @@ +#ifndef FRAMEREADER_HPP +#define FRAMEREADER_HPP + +#include +#include +#include +#include +#include +#include +#include + +#include "channel.hpp" + +// independent of QT, needs ffmpeg +extern "C" { +#include +#include +#include +} + + +class FrameReader { +public: + FrameReader(const char *fn); + uint8_t *get(int idx); + AVFrame *toRGB(AVFrame *); + void waitForReady() { + while (!joined) usleep(10*1000); + } + int getRGBSize() { return width*height*3; } + void loaderThread(); + void cacherThread(); +private: + AVFormatContext *pFormatCtx = NULL; + AVCodecContext *pCodecCtx = NULL; + + struct SwsContext *sws_ctx = NULL; + + int width = 1164; + int height = 874; + + std::vector pkts; + + std::thread *t; + bool joined = false; + + std::map cache; + std::mutex mcache; + + void GOPCache(int idx); + channel to_cache; + + bool valid = true; + char url[0x400]; +}; + +#endif + diff --git a/tools/clib/SConscript b/tools/clib/SConscript new file mode 100644 index 0000000000..6f33fe75bc --- /dev/null +++ b/tools/clib/SConscript @@ -0,0 +1,8 @@ +Import('env') +from sysconfig import get_paths +env['CPPPATH'] += [get_paths()['include']] + +from Cython.Build import cythonize +cythonize("cframereader.pyx") +env.SharedLibrary(File('cframereader.so'), ['cframereader.cpp', 'FrameReader.cpp'], LIBS=['avformat', 'avcodec', 'avutil', 'swscale']) + diff --git a/tools/clib/cframereader.pyx b/tools/clib/cframereader.pyx new file mode 100644 index 0000000000..ffbc277587 --- /dev/null +++ b/tools/clib/cframereader.pyx @@ -0,0 +1,20 @@ +# distutils: language = c++ +# cython: language_level=3 + +cdef extern from "FrameReader.hpp": + cdef cppclass CFrameReader "FrameReader": + CFrameReader(const char *) + char *get(int) + +cdef class FrameReader(): + cdef CFrameReader *fr + + def __cinit__(self, fn): + self.fr = new CFrameReader(fn) + + def __dealloc__(self): + del self.fr + + def get(self, idx): + self.fr.get(idx) + diff --git a/tools/clib/channel.hpp b/tools/clib/channel.hpp new file mode 100644 index 0000000000..d1ce657cec --- /dev/null +++ b/tools/clib/channel.hpp @@ -0,0 +1,35 @@ +#ifndef CHANNEL_HPP +#define CHANNEL_HPP + +#include +#include +#include + +template +class channel { +private: + std::list queue; + std::mutex m; + std::condition_variable cv; +public: + void put(const item &i) { + std::unique_lock lock(m); + queue.push_back(i); + cv.notify_one(); + } + void put_front(const item &i) { + std::unique_lock lock(m); + queue.push_front(i); + cv.notify_one(); + } + item get() { + std::unique_lock lock(m); + cv.wait(lock, [&](){ return !queue.empty(); }); + item result = queue.front(); + queue.pop_front(); + return result; + } +}; + +#endif + diff --git a/tools/lib/__init__.py b/tools/lib/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tools/lib/async_generator.py b/tools/lib/async_generator.py new file mode 100644 index 0000000000..ef535f27f5 --- /dev/null +++ b/tools/lib/async_generator.py @@ -0,0 +1,352 @@ +import functools +import threading +import inspect +import sys +import select +import struct +from math import sqrt +from collections import OrderedDict, deque +from time import time + +from tools.lib.pollable_queue import PollableQueue, Empty, Full, ExistentialError + +EndSentinel = object() + + +def _sync_inner_generator(input_queue, *args, **kwargs): + func = args[0] + args = args[1:] + + get = input_queue.get + while True: + item = get() + if item is EndSentinel: + return + + cookie, value = item + yield cookie, func(value, *args, **kwargs) + + +def _async_streamer_async_inner(input_queue, output_queue, generator_func, args, kwargs): + put = output_queue.put + put_end = True + try: + g = generator_func(input_queue, *args, **kwargs) + for item in g: + put((time(), item)) + g.close() + except ExistentialError: + put_end = False + raise + finally: + if put_end: + put((None, EndSentinel)) + +def _running_mean_var(ltc_stats, x): + old_mean, var = ltc_stats + mean = min(600., 0.98 * old_mean + 0.02 * x) + var = min(5., max(0.1, 0.98 * var + 0.02 * (mean - x) * (old_mean - x))) + return mean, var + +def _find_next_resend(sent_messages, ltc_stats): + if not sent_messages: + return None, None + + oldest_sent_idx = sent_messages._OrderedDict__root[1][2] + send_time, _ = sent_messages[oldest_sent_idx] + + # Assume message has been lost if it is >10 standard deviations from mean. + mean, var = ltc_stats + next_resend_time = send_time + mean + 40. * sqrt(var) + + return oldest_sent_idx, next_resend_time + + +def _do_cleanup(input_queue, output_queue, num_workers, sentinels_received, num_outstanding): + input_fd = input_queue.put_fd() + output_fd = output_queue.get_fd() + + poller = select.epoll() + poller.register(input_fd, select.EPOLLOUT) + poller.register(output_fd, select.EPOLLIN) + + remaining_outputs = [] + end_sentinels_to_send = num_workers - sentinels_received + while sentinels_received < num_workers: + evts = dict(poller.poll(-1 if num_outstanding > 0 else 10.)) + if not evts: + # Workers aren't responding, crash. + break + + if output_fd in evts: + _, maybe_sentinel = output_queue.get() + if maybe_sentinel is EndSentinel: + sentinels_received += 1 + else: + remaining_outputs.append(maybe_sentinel[1]) + num_outstanding -= 1 + + if input_fd in evts: + if end_sentinels_to_send > 0: + input_queue.put_nowait(EndSentinel) + end_sentinels_to_send -= 1 + else: + poller.modify(input_fd, 0) + + # TODO: Raise an exception when a queue thread raises one. + assert sentinels_received == num_workers, (sentinels_received, num_workers) + assert output_queue.empty() + return remaining_outputs + +def _generate_results(input_stream, input_queue, worker_output_queue, output_queue, + num_workers, max_outstanding): + pack_cookie = struct.pack + + # Maps idx -> (send_time, input) + sent_messages = OrderedDict() + oldest_sent_idx = None + next_resend_time = None + ltc_stats = 5., 10. + + # Maps idx -> result + received_messages = {} + next_out = 0 + + # Start things off by pulling the first value. + next_in_item = next(input_stream, EndSentinel) + inputs_remain = next_in_item is not EndSentinel + sentinels_received = 0 + + input_fd = input_queue.put_fd() + worker_output_fd = worker_output_queue.get_fd() + output_fd = output_queue.put_fd() + + poller = select.epoll() + poller.register(input_fd, select.EPOLLOUT) + poller.register(worker_output_fd, select.EPOLLIN) + poller.register(output_fd, 0) + + # Keep sending/retrying until the input stream and sent messages are all done. + while sentinels_received < num_workers and (inputs_remain or sent_messages): + if max_outstanding: + can_send_new = (len(sent_messages) < max_outstanding and + len(received_messages) < max_outstanding and inputs_remain) + else: + can_send_new = inputs_remain + + if (next_resend_time and now >= next_resend_time) or can_send_new: + poller.modify(input_fd, select.EPOLLOUT) + else: + poller.modify(input_fd, 0) + + if next_resend_time: + t = max(0, next_resend_time - now) + evts = dict(poller.poll(t)) + else: + evts = dict(poller.poll()) + now = time() + + if output_fd in evts: + output_queue.put_nowait(received_messages.pop(next_out)) + next_out += 1 + + if next_out not in received_messages: + poller.modify(output_fd, 0) + + if worker_output_fd in evts: + for receive_time, maybe_sentinel in worker_output_queue.get_multiple_nowait(): + # Check for EndSentinel in case of worker crash. + if maybe_sentinel is EndSentinel: + sentinels_received += 1 + continue + idx_bytes, value = maybe_sentinel + idx = struct.unpack("= (3,0): + import queue + import pickle + from io import BytesIO as StringIO +else: + import Queue as queue + import cPickle as pickle + from cStringIO import StringIO + +import subprocess +from aenum import Enum +from lru import LRU +from functools import wraps +from concurrent.futures import ThreadPoolExecutor, as_completed + +from tools.lib.cache import cache_path_for_file_path +from tools.lib.exceptions import DataUnreadableError +try: + from xx.chffr.lib.filereader import FileReader +except ImportError: + from tools.lib.filereader import FileReader +from tools.lib.file_helpers import atomic_write_in_dir +from tools.lib.mkvparse import mkvindex +from tools.lib.route import Route + +H264_SLICE_P = 0 +H264_SLICE_B = 1 +H264_SLICE_I = 2 + +HEVC_SLICE_B = 0 +HEVC_SLICE_P = 1 +HEVC_SLICE_I = 2 + +SLICE_I = 2 # hevc and h264 are the same :) + +class FrameType(Enum): + raw = 1 + h265_stream = 2 + h264_mp4 = 3 + h264_pstream = 4 + ffv1_mkv = 5 + ffvhuff_mkv = 6 + +def fingerprint_video(fn): + with FileReader(fn) as f: + header = f.read(4) + if len(header) == 0: + raise DataUnreadableError("%s is empty" % fn) + elif header == b"\x00\xc0\x12\x00": + return FrameType.raw + elif header == b"\x00\x00\x00\x01": + if 'hevc' in fn: + return FrameType.h265_stream + elif os.path.basename(fn) in ("camera", "acamera"): + return FrameType.h264_pstream + else: + raise NotImplementedError(fn) + elif header == b"\x00\x00\x00\x1c": + return FrameType.h264_mp4 + elif header == b"\x1a\x45\xdf\xa3": + return FrameType.ffv1_mkv + else: + raise NotImplementedError(fn) + + +def ffprobe(fn, fmt=None): + cmd = ["ffprobe", + "-v", "quiet", + "-print_format", "json", + "-show_format", "-show_streams"] + if fmt: + cmd += ["-format", fmt] + cmd += [fn] + + try: + ffprobe_output = subprocess.check_output(cmd) + except subprocess.CalledProcessError as e: + raise DataUnreadableError(fn) + + return json.loads(ffprobe_output) + + +def vidindex(fn, typ): + vidindex_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "vidindex") + vidindex = os.path.join(vidindex_dir, "vidindex") + + subprocess.check_call(["make"], cwd=vidindex_dir, stdout=open("/dev/null","w")) + + with tempfile.NamedTemporaryFile() as prefix_f, \ + tempfile.NamedTemporaryFile() as index_f: + try: + subprocess.check_call([vidindex, typ, fn, prefix_f.name, index_f.name]) + except subprocess.CalledProcessError as e: + raise DataUnreadableError("vidindex failed on file %s" % fn) + with open(index_f.name, "rb") as f: + index = f.read() + with open(prefix_f.name, "rb") as f: + prefix = f.read() + + index = np.frombuffer(index, np.uint32).reshape(-1, 2) + + assert index[-1, 0] == 0xFFFFFFFF + assert index[-1, 1] == os.path.getsize(fn) + + return index, prefix + + +def cache_fn(func): + @wraps(func) + def cache_inner(fn, *args, **kwargs): + cache_prefix = kwargs.pop('cache_prefix', None) + cache_path = cache_path_for_file_path(fn, cache_prefix) + + if cache_path and os.path.exists(cache_path): + with open(cache_path, "rb") as cache_file: + cache_value = pickle.load(cache_file) + else: + cache_value = func(fn, *args, **kwargs) + + if cache_path: + with atomic_write_in_dir(cache_path, mode="wb", overwrite=True) as cache_file: + pickle.dump(cache_value, cache_file, -1) + + return cache_value + + return cache_inner + +@cache_fn +def index_stream(fn, typ): + assert typ in ("hevc", "h264") + + with FileReader(fn) as f: + assert os.path.exists(f.name), fn + index, prefix = vidindex(f.name, typ) + probe = ffprobe(f.name, typ) + + return { + 'index': index, + 'global_prefix': prefix, + 'probe': probe + } + +@cache_fn +def index_mp4(fn): + with FileReader(fn) as f: + return vidindex_mp4(f.name) + +@cache_fn +def index_mkv(fn): + with FileReader(fn) as f: + probe = ffprobe(f.name, "matroska") + with open(f.name, "rb") as d_f: + config_record, index = mkvindex.mkvindex(d_f) + return { + 'probe': probe, + 'config_record': config_record, + 'index': index + } + +def index_videos(camera_paths, cache_prefix=None): + """Requires that paths in camera_paths are contiguous and of the same type.""" + if len(camera_paths) < 1: + raise ValueError("must provide at least one video to index") + + frame_type = fingerprint_video(camera_paths[0]) + if frame_type == FrameType.h264_pstream: + index_pstream(camera_paths, "h264", cache_prefix) + else: + for fn in camera_paths: + index_video(fn, frame_type, cache_prefix) + +def index_video(fn, frame_type=None, cache_prefix=None): + cache_path = cache_path_for_file_path(fn, cache_prefix) + + if os.path.exists(cache_path): + return + + if frame_type is None: + frame_type = fingerprint_video(fn[0]) + + if frame_type == FrameType.h264_pstream: + #hack: try to index the whole route now + route = Route.from_file_path(fn) + + camera_paths = route.camera_paths() + if fn not in camera_paths: + raise DataUnreadableError("Not a contiguous route camera file: {}".format(fn)) + + print("no pstream cache for %s, indexing route %s now" % (fn, route.name)) + index_pstream(route.camera_paths(), "h264", cache_prefix) + elif frame_type == FrameType.h265_stream: + index_stream(fn, "hevc", cache_prefix=cache_prefix) + elif frame_type == FrameType.h264_mp4: + index_mp4(fn, cache_prefix=cache_prefix) + +def get_video_index(fn, frame_type, cache_prefix=None): + cache_path = cache_path_for_file_path(fn, cache_prefix) + + if not os.path.exists(cache_path): + index_video(fn, frame_type, cache_prefix) + + if not os.path.exists(cache_path): + return None + with open(cache_path, "rb") as cache_file: + return pickle.load(cache_file) + +def pstream_predecompress(fns, probe, indexes, global_prefix, cache_prefix, multithreaded=False): + assert len(fns) == len(indexes) + out_fns = [cache_path_for_file_path(fn, cache_prefix, extension=".predecom.mkv") for fn in fns] + out_exists = map(os.path.exists, out_fns) + if all(out_exists): + return + + w = probe['streams'][0]['width'] + h = probe['streams'][0]['height'] + + frame_size = w*h*3/2 # yuv420p + + decompress_proc = subprocess.Popen( + ["ffmpeg", + "-threads", "0" if multithreaded else "1", + "-vsync", "0", + "-f", "h264", + "-i", "pipe:0", + "-threads", "0" if multithreaded else "1", + "-f", "rawvideo", + "-pix_fmt", "yuv420p", + "pipe:1"], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=open("/dev/null", "wb")) + + def write_thread(): + for fn in fns: + with FileReader(fn) as f: + decompress_proc.stdin.write(f.read()) + decompress_proc.stdin.close() + + def read_frame(): + frame = None + try: + frame = decompress_proc.stdout.read(frame_size) + except (IOError, ValueError): + pass + if frame is None or frame == "" or len(frame) != frame_size: + raise DataUnreadableError("pre-decompression failed for %s" % fn) + return frame + + t = threading.Thread(target=write_thread) + t.daemon = True + t.start() + + try: + for fn, out_fn, out_exist, index in zip(fns, out_fns, out_exists, indexes): + if out_exist: + for fi in range(index.shape[0]-1): + read_frame() + continue + + with atomic_write_in_dir(out_fn, mode="w+b", overwrite=True) as out_tmp: + compress_proc = subprocess.Popen( + ["ffmpeg", + "-threads", "0" if multithreaded else "1", + "-y", + "-vsync", "0", + "-f", "rawvideo", + "-pix_fmt", "yuv420p", + "-s", "%dx%d" % (w, h), + "-i", "pipe:0", + "-threads", "0" if multithreaded else "1", + "-f", "matroska", + "-vcodec", "ffv1", + "-g", "0", + out_tmp.name], + stdin=subprocess.PIPE, stderr=open("/dev/null", "wb")) + try: + for fi in range(index.shape[0]-1): + frame = read_frame() + compress_proc.stdin.write(frame) + compress_proc.stdin.close() + except: + compress_proc.kill() + raise + + assert compress_proc.wait() == 0 + + cache_path = cache_path_for_file_path(fn, cache_prefix) + with atomic_write_in_dir(cache_path, mode="wb", overwrite=True) as cache_file: + pickle.dump({ + 'predecom': os.path.basename(out_fn), + 'index': index, + 'probe': probe, + 'global_prefix': global_prefix, + }, cache_file, -1) + + except: + decompress_proc.kill() + raise + finally: + t.join() + + rc = decompress_proc.wait() + if rc != 0: + raise DataUnreadableError(fns[0]) + + +def index_pstream(fns, typ, cache_prefix=None): + if typ != "h264": + raise NotImplementedError(typ) + + if not fns: + raise DataUnreadableError("chffr h264 requires contiguous files") + + out_fns = [cache_path_for_file_path(fn, cache_prefix) for fn in fns] + out_exists = map(os.path.exists, out_fns) + if all(out_exists): return + + # load existing index files to avoid re-doing work + existing_indexes = [] + for out_fn, exists in zip(out_fns, out_exists): + existing = None + if exists: + with open(out_fn, "rb") as cache_file: + existing = pickle.load(cache_file) + existing_indexes.append(existing) + + # probe the first file + if existing_indexes[0]: + probe = existing_indexes[0]['probe'] + else: + with FileReader(fns[0]) as f: + probe = ffprobe(f.name, typ) + + global_prefix = None + + # get the video index of all the segments in this stream + indexes = [] + for i, fn in enumerate(fns): + if existing_indexes[i]: + index = existing_indexes[i]['index'] + prefix = existing_indexes[i]['global_prefix'] + else: + with FileReader(fn) as f: + index, prefix = vidindex(f.name, typ) + if i == 0: + # assert prefix + if not prefix: + raise DataUnreadableError("vidindex failed for %s" % fn) + global_prefix = prefix + indexes.append(index) + + assert global_prefix + + if np.sum(indexes[0][:, 0] == H264_SLICE_I) <= 1: + print("pstream %s is unseekable. pre-decompressing all the segments..." % (fns[0])) + pstream_predecompress(fns, probe, indexes, global_prefix, cache_prefix) + return + + # generate what's required to make each segment self-contained + # (the partial GOP from the end of each segments are put asside to add + # to the start of the following segment) + prefix_data = ["" for _ in fns] + prefix_index = [[] for _ in fns] + for i in range(len(fns)-1): + if indexes[i+1][0, 0] == H264_SLICE_I and indexes[i+1][0, 1] <= 1: + # next file happens to start with a i-frame, dont need use this file's end + continue + + index = indexes[i] + if i == 0 and np.sum(index[:, 0] == H264_SLICE_I) <= 1: + raise NotImplementedError("No I-frames in pstream.") + + # find the last GOP in the index + frame_b = len(index)-1 + while frame_b > 0 and index[frame_b, 0] != H264_SLICE_I: + frame_b -= 1 + + assert frame_b >= 0 + assert index[frame_b, 0] == H264_SLICE_I + + end_len = len(index)-frame_b + + with FileReader(fns[i]) as vid: + vid.seek(index[frame_b, 1]) + end_data = vid.read() + + prefix_data[i+1] = end_data + prefix_index[i+1] = index[frame_b:-1] + # indexes[i] = index[:frame_b] + + for i, fn in enumerate(fns): + cache_path = out_fns[i] + + if os.path.exists(cache_path): + continue + + segment_index = { + 'index': indexes[i], + 'global_prefix': global_prefix, + 'probe': probe, + 'prefix_frame_data': prefix_data[i], # data to prefix the first GOP with + 'num_prefix_frames': len(prefix_index[i]), # number of frames to skip in the first GOP + } + + with atomic_write_in_dir(cache_path, mode="wb", overwrite=True) as cache_file: + pickle.dump(segment_index, cache_file, -1) + +def gpu_info(): + ret = [] + for fn in glob.glob("/proc/driver/nvidia/gpus/*/information"): + with open(fn, "r") as f: + dat = f.read() + kvs = dat.strip().split("\n") + kv = {} + for s in kvs: + k, v = s.split(":", 1) + kv[k] = v.strip() + ret.append(kv) + return ret + +def gpu_supports_hevc(gpuinfo): + return ("GTX 10" in gpuinfo['Model'] or "GTX 20" in gpuinfo['Model'] or gpuinfo['Model'] == "Graphics Device") + +def find_hevc_gpu(): + for gpuinfo in gpu_info(): + if gpu_supports_hevc(gpuinfo): + return int(gpuinfo['Device Minor']) + return None + +def _ffmpeg_fcamera_input_for_frame_info(frame_info): + st = time.time() + fn, num, count, cache_prefix = frame_info + + assert fn.endswith('.hevc') + sindex = index_stream(fn, "hevc", cache_prefix=cache_prefix) + index = sindex['index'] + prefix = sindex['global_prefix'] + probe = sindex['probe'] + + frame_e = num + count + frame_b = num + # must start decoding on an i-frame + while index[frame_b, 0] != HEVC_SLICE_I: + frame_b -= 1 + offset_b = index[frame_b, 1] + offset_e = index[frame_e, 1] + assert frame_b <= num < frame_e + skip = num - frame_b + + w = probe['streams'][0]['width'] + h = probe['streams'][0]['height'] + assert (h, w) == (874, 1164) + + st2 = time.time() + with FileReader(fn) as f: + f.seek(offset_b) + input_data = f.read(offset_e - offset_b) + et = time.time() + + get_time = et-st + get_time2 = et-st2 + + if get_time > 10.0: + print("TOOK OVER 10 seconds to fetch %r %f %f" % (frame_info, get_time, get_time2)) + + return prefix, input_data, skip, count + +def _ffmpeg_fcamera_input_for_frame(pair): + cookie, frame_info = pair + try: + return cookie, _ffmpeg_fcamera_input_for_frame_info(frame_info) + except Exception as e: + # Let the caller handle exceptions. + return cookie, e + + +def _feed_ffmpeg_fcamera_input_work_loop(frames, proc_stdin, select_pipe_fd, cookie_queue): + last_prefix = None + """ + with ThreadPoolExecutor(64) as pool: + futures = [] + for f in frames: + futures.append(pool.submit(_ffmpeg_fcamera_input_for_frame, f)) + for f in as_completed(futures): + cookie, data = f.result() + if isinstance(data, Exception): + # Just print exceptions for now. + print(data) + continue + prefix, input_data, skip, count = data + cookie_queue.put((cookie, count)) + + # Write zeros for skipped frames, ones for keep frames. + os.write(select_pipe_fd, b"\x00" * skip + b"\x01" * count) + + if prefix != last_prefix: + proc_stdin.write(prefix) + last_prefix = prefix + + proc_stdin.write(input_data) + """ + num_threads = 64 + for cookie, data in async_generator( + num_threads, 8 * num_threads, 8 * num_threads, + reliable=False)(_ffmpeg_fcamera_input_for_frame)(frames): + if isinstance(data, Exception): + # Just print exceptions for now. + print(data) + continue + prefix, input_data, skip, count = data + cookie_queue.put((cookie, count)) + + # Write zeros for skipped frames, ones for keep frames. + os.write(select_pipe_fd, b"\x00" * skip + b"\x01" * count) + + if prefix != last_prefix: + proc_stdin.write(prefix) + last_prefix = prefix + + proc_stdin.write(input_data) + +_FCAMERA_FEED_SUCCESS = object() +def feed_ffmpeg_fcamera_input(frames, proc_stdin, select_pipe_fd, cookie_queue): + print("Feed started on {}".format(threading.current_thread().name)) + try: + _feed_ffmpeg_fcamera_input_work_loop(frames, proc_stdin, select_pipe_fd, cookie_queue) + cookie_queue.put((_FCAMERA_FEED_SUCCESS, None)) + finally: + # Always close ffmpeg input. + proc_stdin.close() + + +def read_file_check_size(f, sz, cookie): + buff = bytearray(sz) + bytes_read = f.readinto(buff) + assert bytes_read == sz, (bytes_read, sz) + return buff + + +import signal +import ctypes +def _set_pdeathsig(sig=signal.SIGTERM): + def f(): + libc = ctypes.CDLL('libc.so.6') + return libc.prctl(1, sig) + return f + +def vidindex_mp4(fn): + try: + xmls = subprocess.check_output(["MP4Box", fn, "-diso", "-out", "/dev/stdout"]) + except subprocess.CalledProcessError as e: + raise DataUnreadableError(fn) + + tree = ET.fromstring(xmls) + + def parse_content(s): + assert s.startswith("data:application/octet-string,") + return s[len("data:application/octet-string,"):].decode("hex") + + avc_element = tree.find(".//AVCSampleEntryBox") + width = int(avc_element.attrib['Width']) + height = int(avc_element.attrib['Height']) + + sps_element = avc_element.find(".//AVCDecoderConfigurationRecord/SequenceParameterSet") + pps_element = avc_element.find(".//AVCDecoderConfigurationRecord/PictureParameterSet") + + sps = parse_content(sps_element.attrib['content']) + pps = parse_content(pps_element.attrib['content']) + + media_header = tree.find("MovieBox/TrackBox/MediaBox/MediaHeaderBox") + time_scale = int(media_header.attrib['TimeScale']) + + sample_sizes = [ + int(entry.attrib['Size']) for entry in tree.findall( + "MovieBox/TrackBox/MediaBox/MediaInformationBox/SampleTableBox/SampleSizeBox/SampleSizeEntry") + ] + + sample_dependency = [ + entry.attrib['dependsOnOther'] == "yes" for entry in tree.findall( + "MovieBox/TrackBox/MediaBox/MediaInformationBox/SampleTableBox/SampleDependencyTypeBox/SampleDependencyEntry") + ] + + assert len(sample_sizes) == len(sample_dependency) + + chunk_offsets = [ + int(entry.attrib['offset']) for entry in tree.findall( + "MovieBox/TrackBox/MediaBox/MediaInformationBox/SampleTableBox/ChunkOffsetBox/ChunkEntry") + ] + + sample_chunk_table = [ + (int(entry.attrib['FirstChunk'])-1, int(entry.attrib['SamplesPerChunk'])) for entry in tree.findall( + "MovieBox/TrackBox/MediaBox/MediaInformationBox/SampleTableBox/SampleToChunkBox/SampleToChunkEntry") + ] + + sample_offsets = [None for _ in sample_sizes] + + sample_i = 0 + for i, (first_chunk, samples_per_chunk) in enumerate(sample_chunk_table): + if i == len(sample_chunk_table)-1: + last_chunk = len(chunk_offsets)-1 + else: + last_chunk = sample_chunk_table[i+1][0]-1 + for k in range(first_chunk, last_chunk+1): + sample_offset = chunk_offsets[k] + for _ in range(samples_per_chunk): + sample_offsets[sample_i] = sample_offset + sample_offset += sample_sizes[sample_i] + sample_i += 1 + + assert sample_i == len(sample_sizes) + + pts_offset_table = [ + ( int(entry.attrib['CompositionOffset']), int(entry.attrib['SampleCount']) ) for entry in tree.findall( + "MovieBox/TrackBox/MediaBox/MediaInformationBox/SampleTableBox/CompositionOffsetBox/CompositionOffsetEntry") + ] + sample_pts_offset = [0 for _ in sample_sizes] + sample_i = 0 + for dt, count in pts_offset_table: + for _ in range(count): + sample_pts_offset[sample_i] = dt + sample_i += 1 + + sample_time_table = [ + ( int(entry.attrib['SampleDelta']), int(entry.attrib['SampleCount']) ) for entry in tree.findall( + "MovieBox/TrackBox/MediaBox/MediaInformationBox/SampleTableBox/TimeToSampleBox/TimeToSampleEntry") + ] + sample_time = [None for _ in sample_sizes] + cur_ts = 0 + sample_i = 0 + for dt, count in sample_time_table: + for _ in range(count): + sample_time[sample_i] = (cur_ts + sample_pts_offset[sample_i]) * 1000 / time_scale + + cur_ts += dt + sample_i += 1 + + sample_time.sort() # because we ony decode GOPs in PTS order + + return { + 'width': width, + 'height': height, + 'sample_offsets': sample_offsets, + 'sample_sizes': sample_sizes, + 'sample_dependency': sample_dependency, + 'sample_time': sample_time, + 'sps': sps, + 'pps': pps + } + + +class BaseFrameReader(object): + # properties: frame_type, frame_count, w, h + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + def close(self): + pass + + def get(self, num, count=1, pix_fmt="yuv420p"): + raise NotImplementedError + +def FrameReader(fn, cache_prefix=None, readahead=False, readbehind=False, multithreaded=True): + frame_type = fingerprint_video(fn) + if frame_type == FrameType.raw: + return RawFrameReader(fn) + elif frame_type in (FrameType.h265_stream, FrameType.h264_pstream): + index_data = get_video_index(fn, frame_type, cache_prefix) + if index_data is not None and "predecom" in index_data: + cache_path = cache_path_for_file_path(fn, cache_prefix) + return MKVFrameReader( + os.path.join(os.path.dirname(cache_path), index_data["predecom"])) + else: + return StreamFrameReader(fn, frame_type, index_data, + readahead=readahead, readbehind=readbehind, multithreaded=multithreaded) + elif frame_type == FrameType.h264_mp4: + return MP4FrameReader(fn, readahead=readahead) + elif frame_type == FrameType.ffv1_mkv: + return MKVFrameReader(fn) + else: + raise NotImplementedError(frame_type) + +def rgb24toyuv420(rgb): + yuv_from_rgb = np.array([[ 0.299 , 0.587 , 0.114 ], + [-0.14714119, -0.28886916, 0.43601035 ], + [ 0.61497538, -0.51496512, -0.10001026 ]]) + img = np.dot(rgb.reshape(-1, 3), yuv_from_rgb.T).reshape(rgb.shape) + + y_len = img.shape[0] * img.shape[1] + uv_len = y_len / 4 + + ys = img[:, :, 0] + us = (img[::2, ::2, 1] + img[1::2, ::2, 1] + img[::2, 1::2, 1] + img[1::2, 1::2, 1]) / 4 + 128 + vs = (img[::2, ::2, 2] + img[1::2, ::2, 2] + img[::2, 1::2, 2] + img[1::2, 1::2, 2]) / 4 + 128 + + yuv420 = np.empty(y_len + 2 * uv_len, dtype=img.dtype) + yuv420[:y_len] = ys.reshape(-1) + yuv420[y_len:y_len + uv_len] = us.reshape(-1) + yuv420[y_len + uv_len:y_len + 2 * uv_len] = vs.reshape(-1) + + return yuv420.clip(0,255).astype('uint8') + +class RawData(object): + 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) + +class RawFrameReader(BaseFrameReader): + def __init__(self, fn): + # raw camera + self.fn = fn + self.frame_type = FrameType.raw + self.rawfile = RawData(self.fn) + self.frame_count = self.rawfile.count + self.w, self.h = 640, 480 + + def load_and_debayer(self, img): + img = np.frombuffer(img, dtype='uint8').reshape(960, 1280) + cimg = np.dstack([img[0::2, 1::2], ( + (img[0::2, 0::2].astype("uint16") + img[1::2, 1::2].astype("uint16")) + >> 1).astype("uint8"), img[1::2, 0::2]]) + return cimg + + + def get(self, num, count=1, pix_fmt="yuv420p"): + assert self.frame_count is not None + assert num+count <= self.frame_count + + if pix_fmt not in ("yuv420p", "rgb24"): + raise ValueError("Unsupported pixel format %r" % pix_fmt) + + app = [] + for i in range(num, num+count): + dat = self.rawfile.read(i) + rgb_dat = self.load_and_debayer(dat) + if pix_fmt == "rgb24": + app.append(rgb_dat) + elif pix_fmt == "yuv420p": + app.append(rgb24toyuv420(rgb_dat)) + else: + raise NotImplementedError + + return app + +def decompress_video_data(rawdat, vid_fmt, w, h, pix_fmt, multithreaded=False): + # using a tempfile is much faster than proc.communicate for some reason + + with tempfile.TemporaryFile() as tmpf: + tmpf.write(rawdat) + tmpf.seek(0) + + proc = subprocess.Popen( + ["ffmpeg", + "-threads", "0" if multithreaded else "1", + "-vsync", "0", + "-f", vid_fmt, + "-flags2", "showall", + "-i", "pipe:0", + "-threads", "0" if multithreaded else "1", + "-f", "rawvideo", + "-pix_fmt", pix_fmt, + "pipe:1"], + stdin=tmpf, stdout=subprocess.PIPE, stderr=open("/dev/null")) + + # dat = proc.communicate()[0] + dat = proc.stdout.read() + if proc.wait() != 0: + raise DataUnreadableError("ffmpeg failed") + + if pix_fmt == "rgb24": + ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, h, w, 3) + elif pix_fmt == "yuv420p": + ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, (h*w*3//2)) + elif pix_fmt == "yuv444p": + ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, 3, h, w) + else: + raise NotImplementedError + + return ret + +class VideoStreamDecompressor(object): + def __init__(self, vid_fmt, w, h, pix_fmt, multithreaded=False): + self.vid_fmt = vid_fmt + self.w = w + self.h = h + self.pix_fmt = pix_fmt + + if pix_fmt == "yuv420p": + self.out_size = w*h*3//2 # yuv420p + elif pix_fmt in ("rgb24", "yuv444p"): + self.out_size = w*h*3 + else: + raise NotImplementedError + + self.out_q = queue.Queue() + + self.proc = subprocess.Popen( + ["ffmpeg", + "-threads", "0" if multithreaded else "1", + # "-avioflags", "direct", + "-analyzeduration", "0", + "-probesize", "32", + "-flush_packets", "0", + # "-fflags", "nobuffer", + "-vsync", "0", + "-f", vid_fmt, + "-i", "pipe:0", + "-threads", "0" if multithreaded else "1", + "-f", "rawvideo", + "-pix_fmt", pix_fmt, + "pipe:1"], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=open("/dev/null", "wb")) + + def read_thread(): + while True: + r = self.proc.stdout.read(self.out_size) + if len(r) == 0: + break + assert len(r) == self.out_size + self.out_q.put(r) + + self.t = threading.Thread(target=read_thread) + self.t.daemon = True + self.t.start() + + def __enter__(self): + return self + + def __exit__(self, *args): + self.close() + + def write(self, rawdat): + self.proc.stdin.write(rawdat) + self.proc.stdin.flush() + + def read(self): + dat = self.out_q.get(block=True) + + if self.pix_fmt == "rgb24": + ret = np.frombuffer(dat, dtype=np.uint8).reshape((self.h, self.w, 3)) + elif self.pix_fmt == "yuv420p": + ret = np.frombuffer(dat, dtype=np.uint8) + elif self.pix_fmt == "yuv444p": + ret = np.frombuffer(dat, dtype=np.uint8).reshape((3, self.h, self.w)) + else: + assert False + + return ret + + def eos(self): + self.proc.stdin.close() + + def close(self): + self.proc.stdin.close() + self.t.join() + self.proc.wait() + assert self.proc.wait() == 0 + + +class MKVFrameReader(BaseFrameReader): + def __init__(self, fn): + self.fn = fn + + #print("MKVFrameReader", fn) + index_data = index_mkv(fn) + stream = index_data['probe']['streams'][0] + self.w = stream['width'] + self.h = stream['height'] + + if stream['codec_name'] == 'ffv1': + self.frame_type = FrameType.ffv1_mkv + elif stream['codec_name'] == 'ffvhuff': + self.frame_type = FrameType.ffvhuff_mkv + else: + raise NotImplementedError + + self.config_record = index_data['config_record'] + self.index = index_data['index'] + + self.frame_count = len(self.index) + + def get(self, num, count=1, pix_fmt="yuv420p"): + assert 0 < num+count <= self.frame_count + + frame_dats = [] + with FileReader(self.fn) as f: + for i in range(num, num+count): + pos, length, _ = self.index[i] + f.seek(pos) + frame_dats.append(f.read(length)) + + of = StringIO() + mkvindex.simple_gen(of, self.config_record, self.w, self.h, frame_dats) + + r = decompress_video_data(of.getvalue(), "matroska", self.w, self.h, pix_fmt) + assert len(r) == count + + return r + + +class GOPReader(object): + def get_gop(self, num): + # returns (start_frame_num, num_frames, frames_to_skip, gop_data) + raise NotImplementedError + + +class DoNothingContextManager(object): + def __enter__(self): return self + def __exit__(*x): pass + + +class GOPFrameReader(BaseFrameReader): + #FrameReader with caching and readahead for formats that are group-of-picture based + + def __init__(self, readahead=False, readbehind=False, multithreaded=True): + self.open_ = True + + self.multithreaded = multithreaded + self.readahead = readahead + self.readbehind = readbehind + self.frame_cache = LRU(64) + + if self.readahead: + self.cache_lock = threading.RLock() + self.readahead_last = None + self.readahead_len = 30 + self.readahead_c = threading.Condition() + self.readahead_thread = threading.Thread(target=self._readahead_thread) + self.readahead_thread.daemon = True + self.readahead_thread.start() + else: + self.cache_lock = DoNothingContextManager() + + def close(self): + if not self.open_: + return + self.open_ = False + + if self.readahead: + self.readahead_c.acquire() + self.readahead_c.notify() + self.readahead_c.release() + self.readahead_thread.join() + + def _readahead_thread(self): + while True: + self.readahead_c.acquire() + try: + if not self.open_: + break + self.readahead_c.wait() + finally: + self.readahead_c.release() + if not self.open_: + break + assert self.readahead_last + num, pix_fmt = self.readahead_last + + if self.readbehind: + for k in range(num-1, max(0, num-self.readahead_len), -1): + self._get_one(k, pix_fmt) + else: + for k in range(num, min(self.frame_count, num+self.readahead_len)): + self._get_one(k, pix_fmt) + + def _get_one(self, num, pix_fmt): + assert num < self.frame_count + + if (num, pix_fmt) in self.frame_cache: + return self.frame_cache[(num, pix_fmt)] + + with self.cache_lock: + if (num, pix_fmt) in self.frame_cache: + return self.frame_cache[(num, pix_fmt)] + + frame_b, num_frames, skip_frames, rawdat = self.get_gop(num) + + ret = decompress_video_data(rawdat, self.vid_fmt, self.w, self.h, pix_fmt, + multithreaded=self.multithreaded) + ret = ret[skip_frames:] + assert ret.shape[0] == num_frames + + for i in range(ret.shape[0]): + self.frame_cache[(frame_b+i, pix_fmt)] = ret[i] + + return self.frame_cache[(num, pix_fmt)] + + def get(self, num, count=1, pix_fmt="yuv420p"): + assert self.frame_count is not None + + if num + count > self.frame_count: + raise ValueError("{} > {}".format(num + count, self.frame_count)) + + if pix_fmt not in ("yuv420p", "rgb24", "yuv444p"): + raise ValueError("Unsupported pixel format %r" % pix_fmt) + + ret = [self._get_one(num + i, pix_fmt) for i in range(count)] + + if self.readahead: + self.readahead_last = (num+count, pix_fmt) + self.readahead_c.acquire() + self.readahead_c.notify() + self.readahead_c.release() + + return ret + +class MP4GOPReader(GOPReader): + def __init__(self, fn): + self.fn = fn + self.frame_type = FrameType.h264_mp4 + + self.index = index_mp4(fn) + + self.w = self.index['width'] + self.h = self.index['height'] + self.sample_sizes = self.index['sample_sizes'] + self.sample_offsets = self.index['sample_offsets'] + self.sample_dependency = self.index['sample_dependency'] + + self.vid_fmt = "h264" + + self.frame_count = len(self.sample_sizes) + + self.prefix = "\x00\x00\x00\x01"+self.index['sps']+"\x00\x00\x00\x01"+self.index['pps'] + + def _lookup_gop(self, num): + frame_b = num + while frame_b > 0 and self.sample_dependency[frame_b]: + frame_b -= 1 + + frame_e = num+1 + while frame_e < (len(self.sample_dependency)-1) and self.sample_dependency[frame_e]: + frame_e += 1 + + return (frame_b, frame_e) + + def get_gop(self, num): + frame_b, frame_e = self._lookup_gop(num) + assert frame_b <= num < frame_e + + num_frames = frame_e-frame_b + + with FileReader(self.fn) as f: + rawdat = [] + + sample_i = frame_b + while sample_i < frame_e: + size = self.sample_sizes[sample_i] + start_offset = self.sample_offsets[sample_i] + + # try to read contiguously because a read could actually be a http request + sample_i += 1 + while sample_i < frame_e and size < 10000000 and start_offset+size == self.sample_offsets[sample_i]: + size += self.sample_sizes[sample_i] + sample_i += 1 + + f.seek(start_offset) + sampledat = f.read(size) + + # read length-prefixed NALUs and output in Annex-B + i = 0 + while i < len(sampledat): + nal_len, = struct.unpack(">I", sampledat[i:i+4]) + rawdat.append("\x00\x00\x00\x01"+sampledat[i+4:i+4+nal_len]) + i = i+4+nal_len + assert i == len(sampledat) + + rawdat = self.prefix+''.join(rawdat) + + return frame_b, num_frames, 0, rawdat + +class MP4FrameReader(MP4GOPReader, GOPFrameReader): + def __init__(self, fn, readahead=False): + MP4GOPReader.__init__(self, fn) + GOPFrameReader.__init__(self, readahead) + +class StreamGOPReader(GOPReader): + def __init__(self, fn, frame_type, index_data): + self.fn = fn + + self.frame_type = frame_type + self.frame_count = None + self.w, self.h = None, None + + self.prefix = None + self.index = None + + self.index = index_data['index'] + self.prefix = index_data['global_prefix'] + probe = index_data['probe'] + + if self.frame_type == FrameType.h265_stream: + self.prefix_frame_data = None + self.num_prefix_frames = 0 + self.vid_fmt = "hevc" + + elif self.frame_type == FrameType.h264_pstream: + self.prefix_frame_data = index_data['prefix_frame_data'] + self.num_prefix_frames = index_data['num_prefix_frames'] + + self.vid_fmt = "h264" + + i = 0 + while i < self.index.shape[0] and self.index[i, 0] != SLICE_I: + i += 1 + self.first_iframe = i + + if self.frame_type == FrameType.h265_stream: + assert self.first_iframe == 0 + + self.frame_count = len(self.index)-1 + + self.w = probe['streams'][0]['width'] + self.h = probe['streams'][0]['height'] + + + def _lookup_gop(self, num): + frame_b = num + while frame_b > 0 and self.index[frame_b, 0] != SLICE_I: + frame_b -= 1 + + frame_e = num+1 + while frame_e < (len(self.index)-1) and self.index[frame_e, 0] != SLICE_I: + frame_e += 1 + + offset_b = self.index[frame_b, 1] + offset_e = self.index[frame_e, 1] + + return (frame_b, frame_e, offset_b, offset_e) + + def get_gop(self, num): + frame_b, frame_e, offset_b, offset_e = self._lookup_gop(num) + assert frame_b <= num < frame_e + + num_frames = frame_e-frame_b + + with FileReader(self.fn) as f: + f.seek(offset_b) + rawdat = f.read(offset_e-offset_b) + + if num < self.first_iframe: + assert self.prefix_frame_data + rawdat = self.prefix_frame_data + rawdat + + rawdat = self.prefix + rawdat + + skip_frames = 0 + if num < self.first_iframe: + skip_frames = self.num_prefix_frames + + return frame_b, num_frames, skip_frames, rawdat + +class StreamFrameReader(StreamGOPReader, GOPFrameReader): + def __init__(self, fn, frame_type, index_data, readahead=False, readbehind=False, multithreaded=False): + StreamGOPReader.__init__(self, fn, frame_type, index_data) + GOPFrameReader.__init__(self, readahead, readbehind, multithreaded) + + + + +def GOPFrameIterator(gop_reader, pix_fmt, multithreaded=True): + # this is really ugly. ill think about how to refactor it when i can think good + + IN_FLIGHT_GOPS = 6 # should be enough that the stream decompressor starts returning data + + with VideoStreamDecompressor( + gop_reader.vid_fmt, gop_reader.w, gop_reader.h, pix_fmt, multithreaded) as dec: + + read_work = [] + + def readthing(): + # print read_work, dec.out_q.qsize() + outf = dec.read() + read_thing = read_work[0] + if read_thing[0] > 0: + read_thing[0] -= 1 + else: + assert read_thing[1] > 0 + yield outf + read_thing[1] -= 1 + + if read_thing[1] == 0: + read_work.pop(0) + + i = 0 + while i < gop_reader.frame_count: + frame_b, num_frames, skip_frames, gop_data = gop_reader.get_gop(i) + dec.write(gop_data) + i += num_frames + read_work.append([skip_frames, num_frames]) + + while len(read_work) >= IN_FLIGHT_GOPS: + for v in readthing(): yield v + + dec.eos() + + while read_work: + for v in readthing(): yield v + + +def FrameIterator(fn, pix_fmt, **kwargs): + fr = FrameReader(fn, **kwargs) + if isinstance(fr, GOPReader): + for v in GOPFrameIterator(fr, pix_fmt, kwargs.get("multithreaded", True)): yield v + else: + for i in range(fr.frame_count): + yield fr.get(i, pix_fmt=pix_fmt)[0] + + +def FrameWriter(ofn, frames, vid_fmt=FrameType.ffvhuff_mkv, pix_fmt="rgb24", framerate=20, multithreaded=False): + if pix_fmt not in ("rgb24", "yuv420p"): + raise NotImplementedError + + if vid_fmt == FrameType.ffv1_mkv: + assert ofn.endswith(".mkv") + vcodec = "ffv1" + elif vid_fmt == FrameType.ffvhuff_mkv: + assert ofn.endswith(".mkv") + vcodec = "ffvhuff" + else: + raise NotImplementedError + + frame_gen = iter(frames) + first_frame = next(frame_gen) + + # assert len(frames) > 1 + if pix_fmt == "rgb24": + h, w = first_frame.shape[:2] + elif pix_fmt == "yuv420p": + w = first_frame.shape[1] + h = 2*first_frame.shape[0]//3 + else: + raise NotImplementedError + + compress_proc = subprocess.Popen( + ["ffmpeg", + "-threads", "0" if multithreaded else "1", + "-y", + "-framerate", str(framerate), + "-vsync", "0", + "-f", "rawvideo", + "-pix_fmt", pix_fmt, + "-s", "%dx%d" % (w, h), + "-i", "pipe:0", + "-threads", "0" if multithreaded else "1", + "-f", "matroska", + "-vcodec", vcodec, + "-g", "0", + ofn], + stdin=subprocess.PIPE, stderr=open("/dev/null", "wb")) + try: + compress_proc.stdin.write(first_frame.tobytes()) + for frame in frame_gen: + compress_proc.stdin.write(frame.tobytes()) + compress_proc.stdin.close() + except: + compress_proc.kill() + raise + + assert compress_proc.wait() == 0 + +if __name__ == "__main__": + fn = "cd:/1c79456b0c90f15a/2017-05-10--08-17-00/2/fcamera.hevc" + f = FrameReader(fn) + # print f.get(0, 1).shape + # print f.get(15, 1).shape + for v in GOPFrameIterator(f, "yuv420p"): + print(v) diff --git a/tools/lib/index_log/.gitignore b/tools/lib/index_log/.gitignore new file mode 100644 index 0000000000..17955f06d8 --- /dev/null +++ b/tools/lib/index_log/.gitignore @@ -0,0 +1 @@ +index_log diff --git a/tools/lib/index_log/Makefile b/tools/lib/index_log/Makefile new file mode 100644 index 0000000000..a0264660f8 --- /dev/null +++ b/tools/lib/index_log/Makefile @@ -0,0 +1,19 @@ +CC := gcc +CXX := g++ +PHONELIBS?=../../../../phonelibs + +UNAME_S := $(shell uname -s) +ifeq ($(UNAME_S),Darwin) + CAPNP_FLAGS := /usr/local/lib/libcapnp.a /usr/local/lib/libkj.a +else + CAPNP_FLAGS := -I $(PHONELIBS)/capnp-cpp/include -I $(PHONELIBS)/capnp-cpp/include + CAPNP_LIBS := -L $(PHONELIBS)/capnp-cpp/x64/lib -L $(PHONELIBS)/capnp-cpp/x64/lib -l:libcapnp.a -l:libkj.a +endif + +index_log: index_log.cc + $(eval $@_TMP := $(shell mktemp)) + $(CXX) -std=gnu++11 -o $($@_TMP) \ + index_log.cc \ + $(CAPNP_FLAGS) \ + $(CAPNP_LIBS) + mv $($@_TMP) $@ diff --git a/tools/lib/index_log/index_log.cc b/tools/lib/index_log/index_log.cc new file mode 100644 index 0000000000..8eb4a5e34e --- /dev/null +++ b/tools/lib/index_log/index_log.cc @@ -0,0 +1,66 @@ +#include +#include +#include +#include + +#include +#include + +#include +#include + +#include +#include + +int main(int argc, char** argv) { + + if (argc != 3) { + printf("usage: %s \n", argv[0]); + return 1; + } + + const std::string log_fn = argv[1]; + const std::string index_fn = argv[2]; + + int log_fd = open(log_fn.c_str(), O_RDONLY, 0); + assert(log_fd >= 0); + + off_t log_size = lseek(log_fd, 0, SEEK_END); + lseek(log_fd, 0, SEEK_SET); + + FILE* index_f = NULL; + if (index_fn == "-") { + index_f = stdout; + } else { + index_f = fopen(index_fn.c_str(), "wb"); + } + assert(index_f); + + void* log_data = mmap(NULL, log_size, PROT_READ, MAP_PRIVATE, log_fd, 0); + assert(log_data); + + auto words = kj::arrayPtr((const capnp::word*)log_data, log_size/sizeof(capnp::word)); + while (words.size() > 0) { + + uint64_t idx = ((uintptr_t)words.begin() - (uintptr_t)log_data); + // printf("%llu - %ld\n", idx, words.size()); + const char* idx_bytes = (const char*)&idx; + fwrite(idx_bytes, 8, 1, index_f); + try { + capnp::FlatArrayMessageReader reader(words); + words = kj::arrayPtr(reader.getEnd(), words.end()); + } catch (kj::Exception exc) { + break; + } + + + } + + munmap(log_data, log_size); + + fclose(index_f); + + close(log_fd); + + return 0; +} diff --git a/tools/lib/kbhit.py b/tools/lib/kbhit.py new file mode 100644 index 0000000000..30587590a4 --- /dev/null +++ b/tools/lib/kbhit.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python +import os +import sys +import termios +import atexit +from select import select + +class KBHit: + + def __init__(self): + '''Creates a KBHit object that you can call to do various keyboard things. + ''' + + self.set_kbhit_terminal() + + def set_kbhit_terminal(self): + # Save the terminal settings + self.fd = sys.stdin.fileno() + self.new_term = termios.tcgetattr(self.fd) + self.old_term = termios.tcgetattr(self.fd) + + # New terminal setting unbuffered + self.new_term[3] = (self.new_term[3] & ~termios.ICANON & ~termios.ECHO) + termios.tcsetattr(self.fd, termios.TCSAFLUSH, self.new_term) + + # Support normal-terminal reset at exit + atexit.register(self.set_normal_term) + + def set_normal_term(self): + ''' Resets to normal terminal. On Windows this is a no-op. + ''' + + if os.name == 'nt': + pass + + else: + termios.tcsetattr(self.fd, termios.TCSAFLUSH, self.old_term) + + + def getch(self): + ''' Returns a keyboard character after kbhit() has been called. + Should not be called in the same program as getarrow(). + ''' + return sys.stdin.read(1) + + + def getarrow(self): + ''' Returns an arrow-key code after kbhit() has been called. Codes are + 0 : up + 1 : right + 2 : down + 3 : left + Should not be called in the same program as getch(). + ''' + + if os.name == 'nt': + msvcrt.getch() # skip 0xE0 + c = msvcrt.getch() + vals = [72, 77, 80, 75] + + else: + c = sys.stdin.read(3)[2] + vals = [65, 67, 66, 68] + + return vals.index(ord(c.decode('utf-8'))) + + + def kbhit(self): + ''' Returns True if keyboard character was hit, False otherwise. + ''' + dr,dw,de = select([sys.stdin], [], [], 0) + return dr != [] + + +# Test +if __name__ == "__main__": + + kb = KBHit() + + print('Hit any key, or ESC to exit') + + while True: + + if kb.kbhit(): + c = kb.getch() + if ord(c) == 27: # ESC + break + print(c) + + kb.set_normal_term() + + diff --git a/tools/lib/lazy_property.py b/tools/lib/lazy_property.py new file mode 100644 index 0000000000..85c038f287 --- /dev/null +++ b/tools/lib/lazy_property.py @@ -0,0 +1,12 @@ +class lazy_property(object): + """Defines a property whose value will be computed only once and as needed. + + This can only be used on instance methods. + """ + def __init__(self, func): + self._func = func + + def __get__(self, obj_self, cls): + value = self._func(obj_self) + setattr(obj_self, self._func.__name__, value) + return value diff --git a/tools/lib/log_util.py b/tools/lib/log_util.py new file mode 100755 index 0000000000..50641c3e9e --- /dev/null +++ b/tools/lib/log_util.py @@ -0,0 +1,111 @@ +from cereal import log as capnp_log + +def write_can_to_msg(data, src, msg): + if not isinstance(data[0], Sequence): + data = [data] + + can_msgs = msg.init('can', len(data)) + for i, d in enumerate(data): + if d[0] < 0: continue # ios bug + cc = can_msgs[i] + cc.address = d[0] + cc.busTime = 0 + cc.dat = hex_to_str(d[2]) + if len(d) == 4: + cc.src = d[3] + cc.busTime = d[1] + else: + cc.src = src + +def convert_old_pkt_to_new(old_pkt): + m, d = old_pkt + msg = capnp_log.Event.new_message() + + if len(m) == 3: + _, pid, t = m + msg.logMonoTime = t + else: + t, pid = m + msg.logMonoTime = int(t * 1e9) + + last_velodyne_time = None + + if pid == PID_OBD: + write_can_to_msg(d, 0, msg) + elif pid == PID_CAM: + frame = msg.init('frame') + frame.frameId = d[0] + frame.timestampEof = msg.logMonoTime + # iOS + elif pid == PID_IGPS: + loc = msg.init('gpsLocation') + loc.latitude = d[0] + loc.longitude = d[1] + loc.speed = d[2] + loc.timestamp = int(m[0]*1000.0) # on iOS, first number is wall time in seconds + loc.flags = 1 | 4 # has latitude, longitude, and speed. + elif pid == PID_IMOTION: + user_acceleration = d[:3] + gravity = d[3:6] + + # iOS separates gravity from linear acceleration, so we recombine them. + # Apple appears to use this constant for the conversion. + g = -9.8 + acceleration = [g*(a + b) for a, b in zip(user_acceleration, gravity)] + + accel_event = msg.init('sensorEvents', 1)[0] + accel_event.acceleration.v = acceleration + # android + elif pid == PID_GPS: + if len(d) <= 6 or d[-1] == "gps": + loc = msg.init('gpsLocation') + loc.latitude = d[0] + loc.longitude = d[1] + loc.speed = d[2] + if len(d) > 6: + loc.timestamp = d[6] + loc.flags = 1 | 4 # has latitude, longitude, and speed. + elif pid == PID_ACCEL: + val = d[2] if type(d[2]) != type(0.0) else d + accel_event = msg.init('sensorEvents', 1)[0] + accel_event.acceleration.v = val + elif pid == PID_GYRO: + val = d[2] if type(d[2]) != type(0.0) else d + gyro_event = msg.init('sensorEvents', 1)[0] + gyro_event.init('gyro').v = val + elif pid == PID_LIDAR: + lid = msg.init('lidarPts') + lid.idx = d[3] + elif pid == PID_APPLANIX: + loc = msg.init('liveLocation') + loc.status = d[18] + + loc.lat, loc.lon, loc.alt = d[0:3] + loc.vNED = d[3:6] + + loc.roll = d[6] + loc.pitch = d[7] + loc.heading = d[8] + + loc.wanderAngle = d[9] + loc.trackAngle = d[10] + + loc.speed = d[11] + + loc.gyro = d[12:15] + loc.accel = d[15:18] + elif pid == PID_IBAROMETER: + pressure_event = msg.init('sensorEvents', 1)[0] + _, pressure = d[0:2] + pressure_event.init('pressure').v = [pressure] # Kilopascals + elif pid == PID_IINIT and len(d) == 4: + init_event = msg.init('initData') + init_event.deviceType = capnp_log.InitData.DeviceType.chffrIos + + build_info = init_event.init('iosBuildInfo') + build_info.appVersion = d[0] + build_info.appBuild = int(d[1]) + build_info.osVersion = d[2] + build_info.deviceModel = d[3] + + return msg.as_reader() diff --git a/tools/lib/logreader.py b/tools/lib/logreader.py new file mode 100644 index 0000000000..ebc6ced7ab --- /dev/null +++ b/tools/lib/logreader.py @@ -0,0 +1,205 @@ +import os +import sys +import gzip +import zlib +import json +import bz2 +import tempfile +import requests +import subprocess +from aenum import Enum +import capnp +import numpy as np + +import platform + +from tools.lib.exceptions import DataUnreadableError +try: + from xx.chffr.lib.filereader import FileReader +except ImportError: + from tools.lib.filereader import FileReader +from tools.lib.log_util import convert_old_pkt_to_new +from cereal import log as capnp_log + +OP_PATH = os.path.dirname(os.path.dirname(capnp_log.__file__)) + +def index_log(fn): + index_log_dir = os.path.join(os.path.dirname(os.path.realpath(__file__)), "index_log") + index_log = os.path.join(index_log_dir, "index_log") + phonelibs_dir = os.path.join(OP_PATH, 'phonelibs') + + subprocess.check_call(["make", "PHONELIBS=" + phonelibs_dir], cwd=index_log_dir, stdout=subprocess.DEVNULL) + + try: + dat = subprocess.check_output([index_log, fn, "-"]) + except subprocess.CalledProcessError: + raise DataUnreadableError("%s capnp is corrupted/truncated" % fn) + return np.frombuffer(dat, dtype=np.uint64) + +def event_read_multiple_bytes(dat): + with tempfile.NamedTemporaryFile() as dat_f: + dat_f.write(dat) + dat_f.flush() + idx = index_log(dat_f.name) + + end_idx = np.uint64(len(dat)) + idx = np.append(idx, end_idx) + + return [capnp_log.Event.from_bytes(dat[idx[i]:idx[i+1]]) + for i in range(len(idx)-1)] + + +# this is an iterator itself, and uses private variables from LogReader +class MultiLogIterator(object): + def __init__(self, log_paths, wraparound=True): + self._log_paths = log_paths + self._wraparound = wraparound + + self._first_log_idx = next(i for i in range(len(log_paths)) if log_paths[i] is not None) + self._current_log = self._first_log_idx + self._idx = 0 + self._log_readers = [None]*len(log_paths) + self.start_time = self._log_reader(self._first_log_idx)._ts[0] + + def _log_reader(self, i): + if self._log_readers[i] is None and self._log_paths[i] is not None: + log_path = self._log_paths[i] + print("LogReader:%s" % log_path) + self._log_readers[i] = LogReader(log_path) + + return self._log_readers[i] + + def __iter__(self): + return self + + def _inc(self): + lr = self._log_reader(self._current_log) + if self._idx < len(lr._ents)-1: + self._idx += 1 + else: + self._idx = 0 + self._current_log = next(i for i in range(self._current_log + 1, len(self._log_readers) + 1) if i == len(self._log_readers) or self._log_paths[i] is not None) + # wraparound + if self._current_log == len(self._log_readers): + if self._wraparound: + self._current_log = self._first_log_idx + else: + raise StopIteration + + def __next__(self): + while 1: + lr = self._log_reader(self._current_log) + ret = lr._ents[self._idx] + if lr._do_conversion: + ret = convert_old_pkt_to_new(ret, lr.data_version) + self._inc() + return ret + + def tell(self): + # returns seconds from start of log + return (self._log_reader(self._current_log)._ts[self._idx] - self.start_time) * 1e-9 + + def seek(self, ts): + # seek to nearest minute + minute = int(ts/60) + if minute >= len(self._log_paths) or self._log_paths[minute] is None: + return False + + self._current_log = minute + + # HACK: O(n) seek afterward + self._idx = 0 + while self.tell() < ts: + self._inc() + return True + + +class LogReader(object): + def __init__(self, fn, canonicalize=True, only_union_types=False): + _, ext = os.path.splitext(fn) + data_version = None + + with FileReader(fn) as f: + dat = f.read() + + # decompress file + if ext == ".gz" and ("log_" in fn or "log2" in fn): + dat = zlib.decompress(dat, zlib.MAX_WBITS|32) + elif ext == ".bz2": + dat = bz2.decompress(dat) + elif ext == ".7z": + if platform.system() == "Darwin": + os.environ["LA_LIBRARY_FILEPATH"] = "/usr/local/opt/libarchive/lib/libarchive.dylib" + import libarchive.public + with libarchive.public.memory_reader(dat) as aa: + mdat = [] + for it in aa: + for bb in it.get_blocks(): + mdat.append(bb) + dat = ''.join(mdat) + + # TODO: extension shouln't be a proxy for DeviceType + if ext == "": + if dat[0] == "[": + needs_conversion = True + ents = [json.loads(x) for x in dat.strip().split("\n")[:-1]] + if "_" in fn: + data_version = fn.split("_")[1] + else: + # old rlogs weren't bz2 compressed + needs_conversion = False + ents = event_read_multiple_bytes(dat) + elif ext == ".gz": + if "log_" in fn: + # Zero data file. + ents = [json.loads(x) for x in dat.strip().split("\n")[:-1]] + needs_conversion = True + elif "log2" in fn: + needs_conversion = False + ents = event_read_multiple_bytes(dat) + else: + raise Exception("unknown extension") + elif ext == ".bz2": + needs_conversion = False + ents = event_read_multiple_bytes(dat) + elif ext == ".7z": + needs_conversion = True + ents = [json.loads(x) for x in dat.strip().split("\n")] + else: + raise Exception("unknown extension") + + if needs_conversion: + # TODO: should we call convert_old_pkt_to_new to generate this? + self._ts = [x[0][0]*1e9 for x in ents] + else: + self._ts = [x.logMonoTime for x in ents] + + self.data_version = data_version + self._do_conversion = needs_conversion and canonicalize + self._only_union_types = only_union_types + self._ents = ents + + def __iter__(self): + for ent in self._ents: + if self._do_conversion: + yield convert_old_pkt_to_new(ent, self.data_version) + elif self._only_union_types: + try: + ent.which() + yield ent + except capnp.lib.capnp.KjException: + pass + else: + yield ent + +def load_many_logs_canonical(log_paths): + """Load all logs for a sequence of log paths.""" + for log_path in log_paths: + for msg in LogReader(log_path): + yield msg + +if __name__ == "__main__": + log_path = sys.argv[1] + lr = LogReader(log_path) + for msg in lr: + print(msg) diff --git a/tools/lib/mkvparse/README.md b/tools/lib/mkvparse/README.md new file mode 100644 index 0000000000..2d82f190b9 --- /dev/null +++ b/tools/lib/mkvparse/README.md @@ -0,0 +1,24 @@ +Simple easy-to-use hacky matroska parser + +Define your handler class: + + class MyMatroskaHandler(mkvparse.MatroskaHandler): + def tracks_available(self): + ... + + def segment_info_available(self): + ... + + def frame(self, track_id, timestamp, data, more_laced_blocks, duration, keyframe_flag, invisible_flag, discardable_flag): + ... + +and `mkvparse.mkvparse(file, MyMatroskaHandler())` + + +Supports lacing and setting global timecode scale, subtitles (BlockGroup). Does not support cues, tags, chapters, seeking and so on. Supports resyncing when something bad is encountered in matroska stream. + +Also contains example of generation of Matroska files from python + +Subtitles should remain as text, binary data gets encoded to hex. + +Licence=MIT diff --git a/tools/lib/mkvparse/__init__.py b/tools/lib/mkvparse/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tools/lib/mkvparse/mkvgen.py b/tools/lib/mkvparse/mkvgen.py new file mode 100755 index 0000000000..963d60d019 --- /dev/null +++ b/tools/lib/mkvparse/mkvgen.py @@ -0,0 +1,188 @@ +#!/usr/bin/env python +import sys +import random +import math + +# Simple hacky Matroska generator +# Reads mp3 file "q.mp3" and jpeg images from img/0.jpg, img/1.jpg and so on and +# writes Matroska file with mjpeg and mp3 to stdout + +# License=MIT + +# unsigned +def big_endian_number(number): + if(number<0x100): + return chr(number) + return big_endian_number(number>>8) + chr(number&0xFF) + +ben=big_endian_number + +def ebml_encode_number(number): + def trailing_bits(rest_of_number, number_of_bits): + # like big_endian_number, but can do padding zeroes + if number_of_bits==8: + return chr(rest_of_number&0xFF); + else: + return trailing_bits(rest_of_number>>8, number_of_bits-8) + chr(rest_of_number&0xFF) + + if number == -1: + return chr(0xFF) + if number < 2**7 - 1: + return chr(number|0x80) + if number < 2**14 - 1: + return chr(0x40 | (number>>8)) + trailing_bits(number, 8) + if number < 2**21 - 1: + return chr(0x20 | (number>>16)) + trailing_bits(number, 16) + if number < 2**28 - 1: + return chr(0x10 | (number>>24)) + trailing_bits(number, 24) + if number < 2**35 - 1: + return chr(0x08 | (number>>32)) + trailing_bits(number, 32) + if number < 2**42 - 1: + return chr(0x04 | (number>>40)) + trailing_bits(number, 40) + if number < 2**49 - 1: + return chr(0x02 | (number>>48)) + trailing_bits(number, 48) + if number < 2**56 - 1: + return chr(0x01) + trailing_bits(number, 56) + raise Exception("NUMBER TOO BIG") + +def ebml_element(element_id, data, length=None): + if length==None: + length = len(data) + return big_endian_number(element_id) + ebml_encode_number(length) + data + + +def write_ebml_header(f, content_type, version, read_version): + f.write( + ebml_element(0x1A45DFA3, "" # EBML + + ebml_element(0x4286, ben(1)) # EBMLVersion + + ebml_element(0x42F7, ben(1)) # EBMLReadVersion + + ebml_element(0x42F2, ben(4)) # EBMLMaxIDLength + + ebml_element(0x42F3, ben(8)) # EBMLMaxSizeLength + + ebml_element(0x4282, content_type) # DocType + + ebml_element(0x4287, ben(version)) # DocTypeVersion + + ebml_element(0x4285, ben(read_version)) # DocTypeReadVersion + )) + +def write_infinite_segment_header(f): + # write segment element header + f.write(ebml_element(0x18538067,"",-1)) # Segment (unknown length) + +def random_uid(): + def rint(): + return int(random.random()*(0x100**4)) + return ben(rint()) + ben(rint()) + ben(rint()) + ben(rint()) + + +def example(): + write_ebml_header(sys.stdout, "matroska", 2, 2) + write_infinite_segment_header(sys.stdout) + + + # write segment info (optional) + sys.stdout.write(ebml_element(0x1549A966, "" # SegmentInfo + + ebml_element(0x73A4, random_uid()) # SegmentUID + + ebml_element(0x7BA9, "mkvgen.py test") # Title + + ebml_element(0x4D80, "mkvgen.py") # MuxingApp + + ebml_element(0x5741, "mkvgen.py") # WritingApp + )) + + # write trans data (codecs etc.) + sys.stdout.write(ebml_element(0x1654AE6B, "" # Tracks + + ebml_element(0xAE, "" # TrackEntry + + ebml_element(0xD7, ben(1)) # TrackNumber + + ebml_element(0x73C5, ben(0x77)) # TrackUID + + ebml_element(0x83, ben(0x01)) # TrackType + #0x01 track is a video track + #0x02 track is an audio track + #0x03 track is a complex track, i.e. a combined video and audio track + #0x10 track is a logo track + #0x11 track is a subtitle track + #0x12 track is a button track + #0x20 track is a control track + + ebml_element(0x536E, "mjpeg data") # Name + + ebml_element(0x86, "V_MJPEG") # CodecID + #+ ebml_element(0x23E383, ben(100000000)) # DefaultDuration (opt.), nanoseconds + #+ ebml_element(0x6DE7, ben(100)) # MinCache + + ebml_element(0xE0, "" # Video + + ebml_element(0xB0, ben(640)) # PixelWidth + + ebml_element(0xBA, ben(480)) # PixelHeight + ) + ) + + ebml_element(0xAE, "" # TrackEntry + + ebml_element(0xD7, ben(2)) # TrackNumber + + ebml_element(0x73C5, ben(0x78)) # TrackUID + + ebml_element(0x83, ben(0x02)) # TrackType + #0x01 track is a video track + #0x02 track is an audio track + #0x03 track is a complex track, i.e. a combined video and audio track + #0x10 track is a logo track + #0x11 track is a subtitle track + #0x12 track is a button track + #0x20 track is a control track + + ebml_element(0x536E, "content of mp3 file") # Name + #+ ebml_element(0x6DE7, ben(100)) # MinCache + + ebml_element(0x86, "A_MPEG/L3") # CodecID + #+ ebml_element(0xE1, "") # Audio + ) + )) + + + mp3file = open("q.mp3", "rb") + mp3file.read(500000); + + def mp3framesgenerator(f): + debt="" + while True: + for i in xrange(0,len(debt)+1): + if i >= len(debt)-1: + debt = debt + f.read(8192) + break + #sys.stderr.write("i="+str(i)+" len="+str(len(debt))+"\n") + if ord(debt[i])==0xFF and (ord(debt[i+1]) & 0xF0)==0XF0 and i>700: + if i>0: + yield debt[0:i] + # sys.stderr.write("len="+str(i)+"\n") + debt = debt[i:] + break + + + mp3 = mp3framesgenerator(mp3file) + mp3.next() + + + for i in xrange(0,530): + framefile = open("img/"+str(i)+".jpg", "rb") + framedata = framefile.read() + framefile.close() + + # write cluster (actual video data) + + if random.random()<1: + sys.stdout.write(ebml_element(0x1F43B675, "" # Cluster + + ebml_element(0xE7, ben(int(i*26*4))) # TimeCode, uint, milliseconds + # + ebml_element(0xA7, ben(0)) # Position, uint + + ebml_element(0xA3, "" # SimpleBlock + + ebml_encode_number(1) # track number + + chr(0x00) + chr(0x00) # timecode, relative to Cluster timecode, sint16, in milliseconds + + chr(0x00) # flags + + framedata + ))) + + for u in xrange(0,4): + mp3f=mp3.next() + if random.random()<1: + sys.stdout.write(ebml_element(0x1F43B675, "" # Cluster + + ebml_element(0xE7, ben(i*26*4+u*26)) # TimeCode, uint, milliseconds + + ebml_element(0xA3, "" # SimpleBlock + + ebml_encode_number(2) # track number + + chr(0x00) + chr(0x00) # timecode, relative to Cluster timecode, sint16, in milliseconds + + chr(0x00) # flags + + mp3f + ))) + + + + + +if __name__ == '__main__': + example() diff --git a/tools/lib/mkvparse/mkvindex.py b/tools/lib/mkvparse/mkvindex.py new file mode 100644 index 0000000000..f985280d9d --- /dev/null +++ b/tools/lib/mkvparse/mkvindex.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python +# Copyright (c) 2016, Comma.ai, Inc. + +import sys +import re +import binascii + +from tools.lib.mkvparse import mkvparse +from tools.lib.mkvparse import mkvgen +from tools.lib.mkvparse.mkvgen import ben, ebml_element, ebml_encode_number + +class MatroskaIndex(mkvparse.MatroskaHandler): + # def __init__(self, banlist, nocluster_mode): + # pass + def __init__(self): + self.frameindex = [] + + def tracks_available(self): + _, self.config_record = self.tracks[1]['CodecPrivate'] + + def frame(self, track_id, timestamp, pos, length, more_laced_frames, duration, + keyframe, invisible, discardable): + self.frameindex.append((pos, length, keyframe)) + + + +def mkvindex(f): + handler = MatroskaIndex() + mkvparse.mkvparse(f, handler) + return handler.config_record, handler.frameindex + + +def simple_gen(of, config_record, w, h, framedata): + mkvgen.write_ebml_header(of, "matroska", 2, 2) + mkvgen.write_infinite_segment_header(of) + + of.write(ebml_element(0x1654AE6B, "" # Tracks + + ebml_element(0xAE, "" # TrackEntry + + ebml_element(0xD7, ben(1)) # TrackNumber + + ebml_element(0x73C5, ben(1)) # TrackUID + + ebml_element(0x83, ben(1)) # TrackType = video track + + ebml_element(0x86, "V_MS/VFW/FOURCC") # CodecID + + ebml_element(0xE0, "" # Video + + ebml_element(0xB0, ben(w)) # PixelWidth + + ebml_element(0xBA, ben(h)) # PixelHeight + ) + + ebml_element(0x63A2, config_record) # CodecPrivate (ffv1 configuration record) + ) + )) + + blocks = [] + for fd in framedata: + blocks.append( + ebml_element(0xA3, "" # SimpleBlock + + ebml_encode_number(1) # track number + + chr(0x00) + chr(0x00) # timecode, relative to Cluster timecode, sint16, in milliseconds + + chr(0x80) # flags (keyframe) + + fd + ) + ) + + of.write(ebml_element(0x1F43B675, "" # Cluster + + ebml_element(0xE7, ben(0)) # TimeCode, uint, milliseconds + # + ebml_element(0xA7, ben(0)) # Position, uint + + ''.join(blocks))) + +if __name__ == "__main__": + import random + + if len(sys.argv) != 2: + print("usage: %s mkvpath" % sys.argv[0]) + with open(sys.argv[1], "rb") as f: + cr, index = mkvindex(f) + + # cr = "280000003002000030010000010018004646563100cb070000000000000000000000000000000000".decode("hex") + + def geti(i): + pos, length = index[i] + with open(sys.argv[1], "rb") as f: + f.seek(pos) + return f.read(length) + + dats = [geti(random.randrange(200)) for _ in xrange(30)] + + with open("tmpout.mkv", "wb") as of: + simple_gen(of, cr, dats) + diff --git a/tools/lib/mkvparse/mkvparse.py b/tools/lib/mkvparse/mkvparse.py new file mode 100644 index 0000000000..114fb4ccea --- /dev/null +++ b/tools/lib/mkvparse/mkvparse.py @@ -0,0 +1,761 @@ +# Licence==MIT; Vitaly "_Vi" Shukela 2012 + +# Simple easy-to-use hacky matroska parser + +# Supports SimpleBlock and BlockGroup, lacing, TimecodeScale. +# Does not support seeking, cues, chapters and other features. +# No proper EOF handling unfortunately + +# See "mkvuser.py" for the example + +import traceback +from struct import unpack + +import sys +import datetime + +if sys.version < '3': + range=xrange +else: + #identity=lambda x:x + def ord(something): + if type(something)==bytes: + if something == b"": + raise StopIteration + return something[0] + else: + return something + +def get_major_bit_number(n): + ''' + Takes uint8, returns number of the most significant bit plus the number with that bit cleared. + Examples: + 0b10010101 -> (0, 0b00010101) + 0b00010101 -> (3, 0b00000101) + 0b01111111 -> (1, 0b00111111) + ''' + if not n: + raise Exception("Bad number") + i=0x80; + r=0 + while not n&i: + r+=1 + i>>=1 + return (r,n&~i); + +def read_matroska_number(f, unmodified=False, signed=False): + ''' + Read ebml number. Unmodified means don't clear the length bit (as in Element IDs) + Returns the number and it's length as a tuple + + See examples in "parse_matroska_number" function + ''' + if unmodified and signed: + raise Exception("Contradictary arguments") + first_byte=f.read(1) + if(first_byte==""): + raise StopIteration + r = ord(first_byte) + (n,r2) = get_major_bit_number(r) + if not unmodified: + r=r2 + # from now "signed" means "negative" + i=n + while i: + r = r * 0x100 + ord(f.read(1)) + i-=1 + if signed: + r-=(2**(7*n+7)-1) + else: + if r==2**(7*n+7)-1: + return (-1, n+1) + return (r,n+1) + +def parse_matroska_number(data, pos, unmodified=False, signed=False): + ''' + Parse ebml number from buffer[pos:]. Just like read_matroska_number. + Unmodified means don't clear the length bit (as in Element IDs) + Returns the number plus the new position in input buffer + + Examples: + "\x81" -> (1, pos+1) + "\x40\x01" -> (1, pos+2) + "\x20\x00\x01" -> (1, pos+3) + "\x3F\xFF\xFF" -> (0x1FFFFF, pos+3) + "\x20\x00\x01" unmodified -> (0x200001, pos+3) + "\xBF" signed -> (0, pos+1) + "\xBE" signed -> (-1, pos+1) + "\xC0" signed -> (1, pos+1) + "\x5F\xEF" signed -> (-16, pos+2) + ''' + if unmodified and signed: + raise Exception("Contradictary arguments") + r = ord(data[pos]) + pos+=1 + (n,r2) = get_major_bit_number(r) + if not unmodified: + r=r2 + # from now "signed" means "negative" + i=n + while i: + r = r * 0x100 + ord(data[pos]) + pos+=1 + i-=1 + if signed: + r-=(2**(7*n+6)-1) + else: + if r==2**(7*n+7)-1: + return (-1, pos) + return (r,pos) + +def parse_xiph_number(data, pos): + ''' + Parse the Xiph lacing number from data[pos:] + Returns the number plus the new position + + Examples: + "\x01" -> (1, pos+1) + "\x55" -> (0x55, pos+1) + "\xFF\x04" -> (0x103, pos+2) + "\xFF\xFF\x04" -> (0x202, pos+3) + "\xFF\xFF\x00" -> (0x1FE, pos+3) + ''' + v = ord(data[pos]) + pos+=1 + + r=0 + while v==255: + r+=v + v = ord(data[pos]) + pos+=1 + + r+=v + return (r, pos) + + +def parse_fixedlength_number(data, pos, length, signed=False): + ''' + Read the big-endian number from data[pos:pos+length] + Returns the number plus the new position + + Examples: + "\x01" -> (0x1, pos+1) + "\x55" -> (0x55, pos+1) + "\x55" signed -> (0x55, pos+1) + "\xFF\x04" -> (0xFF04, pos+2) + "\xFF\x04" signed -> (-0x00FC, pos+2) + ''' + r=0 + for i in range(length): + r=r*0x100+ord(data[pos+i]) + if signed: + if ord(data[pos]) & 0x80: + r-=2**(8*length) + return (r, pos+length) + +def read_fixedlength_number(f, length, signed=False): + """ Read length bytes and parse (parse_fixedlength_number) it. + Returns only the number""" + buf = f.read(length) + (r, pos) = parse_fixedlength_number(buf, 0, length, signed) + return r + +def read_ebml_element_header(f): + ''' + Read Element ID and size + Returns id, element size and this header size + ''' + (id_, n) = read_matroska_number(f, unmodified=True) + (size, n2) = read_matroska_number(f) + return (id_, size, n+n2) + +class EbmlElementType: + VOID=0 + MASTER=1 # read all subelements and return tree. Don't use this too large things like Segment + UNSIGNED=2 + SIGNED=3 + TEXTA=4 + TEXTU=5 + BINARY=6 + FLOAT=7 + DATE=8 + + JUST_GO_ON=10 # For "Segment". + # Actually MASTER, but don't build the tree for all subelements, + # interpreting all child elements as if they were top-level elements + + +EET=EbmlElementType + +# lynx -width=10000 -dump http://matroska.org/technical/specs/index.html +# | sed 's/not 0/not0/g; s/> 0/>0/g; s/Sampling Frequency/SamplingFrequency/g' +# | awk '{print $1 " " $3 " " $8}' +# | grep '\[..\]' +# | perl -ne '/(\S+) (\S+) (.)/; +# $name=$1; $id=$2; $type=$3; +# $id=~s/\[|\]//g; +# %types = (m=>"EET.MASTER", +# u=>"EET.UNSIGNED", +# i=>"EET.SIGNED", +# 8=>"EET.TEXTU", +# s=>"EET.TEXTA", +# b=>"EET.BINARY", +# f=>"EET.FLOAT", +# d=>"EET.DATE"); +# $t=$types{$type}; +# next unless $t; +# $t="EET.JUST_GO_ON" if $name eq "Segment" or $name eq "Cluster"; +# print "\t0x$id: ($t, \"$name\"),\n";' + +element_types_names = { + 0x1A45DFA3: (EET.MASTER, "EBML"), + 0x4286: (EET.UNSIGNED, "EBMLVersion"), + 0x42F7: (EET.UNSIGNED, "EBMLReadVersion"), + 0x42F2: (EET.UNSIGNED, "EBMLMaxIDLength"), + 0x42F3: (EET.UNSIGNED, "EBMLMaxSizeLength"), + 0x4282: (EET.TEXTA, "DocType"), + 0x4287: (EET.UNSIGNED, "DocTypeVersion"), + 0x4285: (EET.UNSIGNED, "DocTypeReadVersion"), + 0xEC: (EET.BINARY, "Void"), + 0xBF: (EET.BINARY, "CRC-32"), + 0x1B538667: (EET.MASTER, "SignatureSlot"), + 0x7E8A: (EET.UNSIGNED, "SignatureAlgo"), + 0x7E9A: (EET.UNSIGNED, "SignatureHash"), + 0x7EA5: (EET.BINARY, "SignaturePublicKey"), + 0x7EB5: (EET.BINARY, "Signature"), + 0x7E5B: (EET.MASTER, "SignatureElements"), + 0x7E7B: (EET.MASTER, "SignatureElementList"), + 0x6532: (EET.BINARY, "SignedElement"), + 0x18538067: (EET.JUST_GO_ON, "Segment"), + 0x114D9B74: (EET.MASTER, "SeekHead"), + 0x4DBB: (EET.MASTER, "Seek"), + 0x53AB: (EET.BINARY, "SeekID"), + 0x53AC: (EET.UNSIGNED, "SeekPosition"), + 0x1549A966: (EET.MASTER, "Info"), + 0x73A4: (EET.BINARY, "SegmentUID"), + 0x7384: (EET.TEXTU, "SegmentFilename"), + 0x3CB923: (EET.BINARY, "PrevUID"), + 0x3C83AB: (EET.TEXTU, "PrevFilename"), + 0x3EB923: (EET.BINARY, "NextUID"), + 0x3E83BB: (EET.TEXTU, "NextFilename"), + 0x4444: (EET.BINARY, "SegmentFamily"), + 0x6924: (EET.MASTER, "ChapterTranslate"), + 0x69FC: (EET.UNSIGNED, "ChapterTranslateEditionUID"), + 0x69BF: (EET.UNSIGNED, "ChapterTranslateCodec"), + 0x69A5: (EET.BINARY, "ChapterTranslateID"), + 0x2AD7B1: (EET.UNSIGNED, "TimecodeScale"), + 0x4489: (EET.FLOAT, "Duration"), + 0x4461: (EET.DATE, "DateUTC"), + 0x7BA9: (EET.TEXTU, "Title"), + 0x4D80: (EET.TEXTU, "MuxingApp"), + 0x5741: (EET.TEXTU, "WritingApp"), + 0x1F43B675: (EET.JUST_GO_ON, "Cluster"), + 0xE7: (EET.UNSIGNED, "Timecode"), + 0x5854: (EET.MASTER, "SilentTracks"), + 0x58D7: (EET.UNSIGNED, "SilentTrackNumber"), + 0xA7: (EET.UNSIGNED, "Position"), + 0xAB: (EET.UNSIGNED, "PrevSize"), + 0xA3: (EET.BINARY, "SimpleBlock"), + 0xA0: (EET.MASTER, "BlockGroup"), + 0xA1: (EET.BINARY, "Block"), + 0xA2: (EET.BINARY, "BlockVirtual"), + 0x75A1: (EET.MASTER, "BlockAdditions"), + 0xA6: (EET.MASTER, "BlockMore"), + 0xEE: (EET.UNSIGNED, "BlockAddID"), + 0xA5: (EET.BINARY, "BlockAdditional"), + 0x9B: (EET.UNSIGNED, "BlockDuration"), + 0xFA: (EET.UNSIGNED, "ReferencePriority"), + 0xFB: (EET.SIGNED, "ReferenceBlock"), + 0xFD: (EET.SIGNED, "ReferenceVirtual"), + 0xA4: (EET.BINARY, "CodecState"), + 0x8E: (EET.MASTER, "Slices"), + 0xE8: (EET.MASTER, "TimeSlice"), + 0xCC: (EET.UNSIGNED, "LaceNumber"), + 0xCD: (EET.UNSIGNED, "FrameNumber"), + 0xCB: (EET.UNSIGNED, "BlockAdditionID"), + 0xCE: (EET.UNSIGNED, "Delay"), + 0xCF: (EET.UNSIGNED, "SliceDuration"), + 0xC8: (EET.MASTER, "ReferenceFrame"), + 0xC9: (EET.UNSIGNED, "ReferenceOffset"), + 0xCA: (EET.UNSIGNED, "ReferenceTimeCode"), + 0xAF: (EET.BINARY, "EncryptedBlock"), + 0x1654AE6B: (EET.MASTER, "Tracks"), + 0xAE: (EET.MASTER, "TrackEntry"), + 0xD7: (EET.UNSIGNED, "TrackNumber"), + 0x73C5: (EET.UNSIGNED, "TrackUID"), + 0x83: (EET.UNSIGNED, "TrackType"), + 0xB9: (EET.UNSIGNED, "FlagEnabled"), + 0x88: (EET.UNSIGNED, "FlagDefault"), + 0x55AA: (EET.UNSIGNED, "FlagForced"), + 0x9C: (EET.UNSIGNED, "FlagLacing"), + 0x6DE7: (EET.UNSIGNED, "MinCache"), + 0x6DF8: (EET.UNSIGNED, "MaxCache"), + 0x23E383: (EET.UNSIGNED, "DefaultDuration"), + 0x23314F: (EET.FLOAT, "TrackTimecodeScale"), + 0x537F: (EET.SIGNED, "TrackOffset"), + 0x55EE: (EET.UNSIGNED, "MaxBlockAdditionID"), + 0x536E: (EET.TEXTU, "Name"), + 0x22B59C: (EET.TEXTA, "Language"), + 0x86: (EET.TEXTA, "CodecID"), + 0x63A2: (EET.BINARY, "CodecPrivate"), + 0x258688: (EET.TEXTU, "CodecName"), + 0x7446: (EET.UNSIGNED, "AttachmentLink"), + 0x3A9697: (EET.TEXTU, "CodecSettings"), + 0x3B4040: (EET.TEXTA, "CodecInfoURL"), + 0x26B240: (EET.TEXTA, "CodecDownloadURL"), + 0xAA: (EET.UNSIGNED, "CodecDecodeAll"), + 0x6FAB: (EET.UNSIGNED, "TrackOverlay"), + 0x6624: (EET.MASTER, "TrackTranslate"), + 0x66FC: (EET.UNSIGNED, "TrackTranslateEditionUID"), + 0x66BF: (EET.UNSIGNED, "TrackTranslateCodec"), + 0x66A5: (EET.BINARY, "TrackTranslateTrackID"), + 0xE0: (EET.MASTER, "Video"), + 0x9A: (EET.UNSIGNED, "FlagInterlaced"), + 0x53B8: (EET.UNSIGNED, "StereoMode"), + 0x53B9: (EET.UNSIGNED, "OldStereoMode"), + 0xB0: (EET.UNSIGNED, "PixelWidth"), + 0xBA: (EET.UNSIGNED, "PixelHeight"), + 0x54AA: (EET.UNSIGNED, "PixelCropBottom"), + 0x54BB: (EET.UNSIGNED, "PixelCropTop"), + 0x54CC: (EET.UNSIGNED, "PixelCropLeft"), + 0x54DD: (EET.UNSIGNED, "PixelCropRight"), + 0x54B0: (EET.UNSIGNED, "DisplayWidth"), + 0x54BA: (EET.UNSIGNED, "DisplayHeight"), + 0x54B2: (EET.UNSIGNED, "DisplayUnit"), + 0x54B3: (EET.UNSIGNED, "AspectRatioType"), + 0x2EB524: (EET.BINARY, "ColourSpace"), + 0x2FB523: (EET.FLOAT, "GammaValue"), + 0x2383E3: (EET.FLOAT, "FrameRate"), + 0xE1: (EET.MASTER, "Audio"), + 0xB5: (EET.FLOAT, "SamplingFrequency"), + 0x78B5: (EET.FLOAT, "OutputSamplingFrequency"), + 0x9F: (EET.UNSIGNED, "Channels"), + 0x7D7B: (EET.BINARY, "ChannelPositions"), + 0x6264: (EET.UNSIGNED, "BitDepth"), + 0xE2: (EET.MASTER, "TrackOperation"), + 0xE3: (EET.MASTER, "TrackCombinePlanes"), + 0xE4: (EET.MASTER, "TrackPlane"), + 0xE5: (EET.UNSIGNED, "TrackPlaneUID"), + 0xE6: (EET.UNSIGNED, "TrackPlaneType"), + 0xE9: (EET.MASTER, "TrackJoinBlocks"), + 0xED: (EET.UNSIGNED, "TrackJoinUID"), + 0xC0: (EET.UNSIGNED, "TrickTrackUID"), + 0xC1: (EET.BINARY, "TrickTrackSegmentUID"), + 0xC6: (EET.UNSIGNED, "TrickTrackFlag"), + 0xC7: (EET.UNSIGNED, "TrickMasterTrackUID"), + 0xC4: (EET.BINARY, "TrickMasterTrackSegmentUID"), + 0x6D80: (EET.MASTER, "ContentEncodings"), + 0x6240: (EET.MASTER, "ContentEncoding"), + 0x5031: (EET.UNSIGNED, "ContentEncodingOrder"), + 0x5032: (EET.UNSIGNED, "ContentEncodingScope"), + 0x5033: (EET.UNSIGNED, "ContentEncodingType"), + 0x5034: (EET.MASTER, "ContentCompression"), + 0x4254: (EET.UNSIGNED, "ContentCompAlgo"), + 0x4255: (EET.BINARY, "ContentCompSettings"), + 0x5035: (EET.MASTER, "ContentEncryption"), + 0x47E1: (EET.UNSIGNED, "ContentEncAlgo"), + 0x47E2: (EET.BINARY, "ContentEncKeyID"), + 0x47E3: (EET.BINARY, "ContentSignature"), + 0x47E4: (EET.BINARY, "ContentSigKeyID"), + 0x47E5: (EET.UNSIGNED, "ContentSigAlgo"), + 0x47E6: (EET.UNSIGNED, "ContentSigHashAlgo"), + 0x1C53BB6B: (EET.MASTER, "Cues"), + 0xBB: (EET.MASTER, "CuePoint"), + 0xB3: (EET.UNSIGNED, "CueTime"), + 0xB7: (EET.MASTER, "CueTrackPositions"), + 0xF7: (EET.UNSIGNED, "CueTrack"), + 0xF1: (EET.UNSIGNED, "CueClusterPosition"), + 0x5378: (EET.UNSIGNED, "CueBlockNumber"), + 0xEA: (EET.UNSIGNED, "CueCodecState"), + 0xDB: (EET.MASTER, "CueReference"), + 0x96: (EET.UNSIGNED, "CueRefTime"), + 0x97: (EET.UNSIGNED, "CueRefCluster"), + 0x535F: (EET.UNSIGNED, "CueRefNumber"), + 0xEB: (EET.UNSIGNED, "CueRefCodecState"), + 0x1941A469: (EET.MASTER, "Attachments"), + 0x61A7: (EET.MASTER, "AttachedFile"), + 0x467E: (EET.TEXTU, "FileDescription"), + 0x466E: (EET.TEXTU, "FileName"), + 0x4660: (EET.TEXTA, "FileMimeType"), + 0x465C: (EET.BINARY, "FileData"), + 0x46AE: (EET.UNSIGNED, "FileUID"), + 0x4675: (EET.BINARY, "FileReferral"), + 0x4661: (EET.UNSIGNED, "FileUsedStartTime"), + 0x4662: (EET.UNSIGNED, "FileUsedEndTime"), + 0x1043A770: (EET.MASTER, "Chapters"), + 0x45B9: (EET.MASTER, "EditionEntry"), + 0x45BC: (EET.UNSIGNED, "EditionUID"), + 0x45BD: (EET.UNSIGNED, "EditionFlagHidden"), + 0x45DB: (EET.UNSIGNED, "EditionFlagDefault"), + 0x45DD: (EET.UNSIGNED, "EditionFlagOrdered"), + 0xB6: (EET.MASTER, "ChapterAtom"), + 0x73C4: (EET.UNSIGNED, "ChapterUID"), + 0x91: (EET.UNSIGNED, "ChapterTimeStart"), + 0x92: (EET.UNSIGNED, "ChapterTimeEnd"), + 0x98: (EET.UNSIGNED, "ChapterFlagHidden"), + 0x4598: (EET.UNSIGNED, "ChapterFlagEnabled"), + 0x6E67: (EET.BINARY, "ChapterSegmentUID"), + 0x6EBC: (EET.UNSIGNED, "ChapterSegmentEditionUID"), + 0x63C3: (EET.UNSIGNED, "ChapterPhysicalEquiv"), + 0x8F: (EET.MASTER, "ChapterTrack"), + 0x89: (EET.UNSIGNED, "ChapterTrackNumber"), + 0x80: (EET.MASTER, "ChapterDisplay"), + 0x85: (EET.TEXTU, "ChapString"), + 0x437C: (EET.TEXTA, "ChapLanguage"), + 0x437E: (EET.TEXTA, "ChapCountry"), + 0x6944: (EET.MASTER, "ChapProcess"), + 0x6955: (EET.UNSIGNED, "ChapProcessCodecID"), + 0x450D: (EET.BINARY, "ChapProcessPrivate"), + 0x6911: (EET.MASTER, "ChapProcessCommand"), + 0x6922: (EET.UNSIGNED, "ChapProcessTime"), + 0x6933: (EET.BINARY, "ChapProcessData"), + 0x1254C367: (EET.MASTER, "Tags"), + 0x7373: (EET.MASTER, "Tag"), + 0x63C0: (EET.MASTER, "Targets"), + 0x68CA: (EET.UNSIGNED, "TargetTypeValue"), + 0x63CA: (EET.TEXTA, "TargetType"), + 0x63C5: (EET.UNSIGNED, "TagTrackUID"), + 0x63C9: (EET.UNSIGNED, "TagEditionUID"), + 0x63C4: (EET.UNSIGNED, "TagChapterUID"), + 0x63C6: (EET.UNSIGNED, "TagAttachmentUID"), + 0x67C8: (EET.MASTER, "SimpleTag"), + 0x45A3: (EET.TEXTU, "TagName"), + 0x447A: (EET.TEXTA, "TagLanguage"), + 0x4484: (EET.UNSIGNED, "TagDefault"), + 0x4487: (EET.TEXTU, "TagString"), + 0x4485: (EET.BINARY, "TagBinary"), + 0x56AA: (EET.UNSIGNED, "CodecDelay"), + 0x56BB: (EET.UNSIGNED, "SeekPreRoll"), + 0xF0: (EET.UNSIGNED, "CueRelativePosition"), + 0x53C0: (EET.UNSIGNED, "AlphaMode"), + 0x55B2: (EET.UNSIGNED, "BitsPerChannel"), + 0x55B5: (EET.UNSIGNED, "CbSubsamplingHorz"), + 0x55B6: (EET.UNSIGNED, "CbSubsamplingVert"), + 0x5654: (EET.TEXTU, "ChapterStringUID"), + 0x55B7: (EET.UNSIGNED, "ChromaSitingHorz"), + 0x55B8: (EET.UNSIGNED, "ChromaSitingVert"), + 0x55B3: (EET.UNSIGNED, "ChromaSubsamplingHorz"), + 0x55B4: (EET.UNSIGNED, "ChromaSubsamplingVert"), + 0x55B0: (EET.MASTER, "Colour"), + 0x234E7A: (EET.UNSIGNED, "DefaultDecodedFieldDuration"), + 0x75A2: (EET.SIGNED, "DiscardPadding"), + 0x9D: (EET.UNSIGNED, "FieldOrder"), + 0x55D9: (EET.FLOAT, "LuminanceMax"), + 0x55DA: (EET.FLOAT, "LuminanceMin"), + 0x55D0: (EET.MASTER, "MasteringMetadata"), + 0x55B1: (EET.UNSIGNED, "MatrixCoefficients"), + 0x55BC: (EET.UNSIGNED, "MaxCLL"), + 0x55BD: (EET.UNSIGNED, "MaxFALL"), + 0x55BB: (EET.UNSIGNED, "Primaries"), + 0x55D5: (EET.FLOAT, "PrimaryBChromaticityX"), + 0x55D6: (EET.FLOAT, "PrimaryBChromaticityY"), + 0x55D3: (EET.FLOAT, "PrimaryGChromaticityX"), + 0x55D4: (EET.FLOAT, "PrimaryGChromaticityY"), + 0x55D1: (EET.FLOAT, "PrimaryRChromaticityX"), + 0x55D2: (EET.FLOAT, "PrimaryRChromaticityY"), + 0x55B9: (EET.UNSIGNED, "Range"), + 0x55BA: (EET.UNSIGNED, "TransferCharacteristics"), + 0x55D7: (EET.FLOAT, "WhitePointChromaticityX"), + 0x55D8: (EET.FLOAT, "WhitePointChromaticityY"), +} + +def read_simple_element(f, type_, size): + date = None + if size==0: + return "" + + if type_==EET.UNSIGNED: + data=read_fixedlength_number(f, size, False) + elif type_==EET.SIGNED: + data=read_fixedlength_number(f, size, True) + elif type_==EET.TEXTA: + data=f.read(size) + data = data.replace(b"\x00", b"") # filter out \0, for gstreamer + data = data.decode("ascii") + elif type_==EET.TEXTU: + data=f.read(size) + data = data.replace(b"\x00", b"") # filter out \0, for gstreamer + data = data.decode("UTF-8") + elif type_==EET.MASTER: + data=read_ebml_element_tree(f, size) + elif type_==EET.DATE: + data=read_fixedlength_number(f, size, True) + data*= 1e-9 + data+= (datetime.datetime(2001, 1, 1) - datetime.datetime(1970, 1, 1)).total_seconds() + # now should be UNIX date + elif type_==EET.FLOAT: + if size==4: + data = f.read(4) + data = unpack(">f", data)[0] + elif size==8: + data = f.read(8) + data = unpack(">d", data)[0] + else: + data=read_fixedlength_number(f, size, False) + sys.stderr.write("mkvparse: Floating point of size %d is not supported\n" % size) + data = None + else: + data=f.read(size) + return data + +def read_ebml_element_tree(f, total_size): + ''' + Build tree of elements, reading f until total_size reached + Don't use for the whole segment, it's not Haskell + + Returns list of pairs (element_name, element_value). + element_value can also be list of pairs + ''' + childs=[] + while(total_size>0): + (id_, size, hsize) = read_ebml_element_header(f) + if size == -1: + sys.stderr.write("mkvparse: Element %x without size? Damaged data? Skipping %d bytes\n" % (id_, size, total_size)) + f.read(total_size); + break; + if size>total_size: + sys.stderr.write("mkvparse: Element %x with size %d? Damaged data? Skipping %d bytes\n" % (id_, size, total_size)) + f.read(total_size); + break + type_ = EET.BINARY + name = "unknown_%x"%id_ + if id_ in element_types_names: + (type_, name) = element_types_names[id_] + data = read_simple_element(f, type_, size) + total_size-=(size+hsize) + childs.append((name, (type_, data))) + return childs + + +class MatroskaHandler: + """ User for mkvparse should override these methods """ + def tracks_available(self): + pass + def segment_info_available(self): + pass + def frame(self, track_id, timestamp, data, more_laced_frames, duration, keyframe, invisible, discardable): + pass + def ebml_top_element(self, id_, name_, type_, data_): + pass + def before_handling_an_element(self): + pass + def begin_handling_ebml_element(self, id_, name, type_, headersize, datasize): + return type_ + def element_data_available(self, id_, name, type_, headersize, data): + pass + +def handle_block(buffer, buffer_pos, handler, cluster_timecode, timecode_scale=1000000, duration=None, header_removal_headers_for_tracks={}): + ''' + Decode a block, handling all lacings, send it to handler with appropriate timestamp, track number + ''' + pos=0 + (tracknum, pos) = parse_matroska_number(buffer, pos, signed=False) + (tcode, pos) = parse_fixedlength_number(buffer, pos, 2, signed=True) + flags = ord(buffer[pos]); pos+=1 + f_keyframe = (flags&0x80 == 0x80) + f_invisible = (flags&0x08 == 0x08) + f_discardable = (flags&0x01 == 0x01) + laceflags=flags&0x06 + + block_timecode = (cluster_timecode + tcode)*(timecode_scale*0.000000001) + + header_removal_prefix = b"" + if tracknum in header_removal_headers_for_tracks: + # header_removal_prefix = header_removal_headers_for_tracks[tracknum] + raise NotImplementedError + + if laceflags == 0x00: # no lacing + # buf = buffer[pos:] + handler.frame(tracknum, block_timecode, buffer_pos+pos, len(buffer)-pos, + 0, duration, f_keyframe, f_invisible, f_discardable) + return + + numframes = ord(buffer[pos]); pos+=1 + numframes+=1 + + lengths=[] + + if laceflags == 0x02: # Xiph lacing + accumlength=0 + for i in range(numframes-1): + (l, pos) = parse_xiph_number(buffer, pos) + lengths.append(l) + accumlength+=l + lengths.append(len(buffer)-pos-accumlength) + elif laceflags == 0x06: # EBML lacing + accumlength=0 + if numframes: + (flength, pos) = parse_matroska_number(buffer, pos, signed=False) + lengths.append(flength) + accumlength+=flength + for i in range(numframes-2): + (l, pos) = parse_matroska_number(buffer, pos, signed=True) + flength+=l + lengths.append(flength) + accumlength+=flength + lengths.append(len(buffer)-pos-accumlength) + elif laceflags==0x04: # Fixed size lacing + fl=int((len(buffer)-pos)/numframes) + for i in range(numframes): + lengths.append(fl) + + more_laced_frames=numframes-1 + for i in lengths: + # buf = buffer[pos:pos+i] + handler.frame(tracknum, block_timecode, buffer_pos+pos, i, more_laced_frames, duration, + f_keyframe, f_invisible, f_discardable) + pos+=i + more_laced_frames-=1 + + +def resync(f): + sys.stderr.write("mvkparse: Resyncing\n") + while True: + b = f.read(1); + if b == b"": return (None, None); + if b == b"\x1F": + b2 = f.read(3); + if b2 == b"\x43\xB6\x75": + (seglen, x) = read_matroska_number(f) + return (0x1F43B675, seglen, x+4) # cluster + if b == b"\x18": + b2 = f.read(3) + if b2 == b"\x53\x80\x67": + (seglen, x) = read_matroska_number(f) + return (0x18538067, seglen, x+4) # segment + if b == b"\x16": + b2 = f.read(3) + if b2 == b"\x54\xAE\x6B": + (seglen ,x )= read_matroska_number(f) + return (0x1654AE6B, seglen, x+4) # tracks + + + + +def mkvparse(f, handler): + ''' + Read mkv file f and call handler methods when track or segment information is ready or when frame is read. + Handles lacing, timecodes (except of per-track scaling) + ''' + timecode_scale = 1000000 + current_cluster_timecode = 0 + resync_element_id = None + resync_element_size = None + resync_element_headersize = None + header_removal_headers_for_tracks = {} + while f: + (id_, size, hsize) = (None, None, None) + tree = None + data = None + (type_, name) = (None, None) + try: + if not resync_element_id: + try: + handler.before_handling_an_element() + (id_, size, hsize) = read_ebml_element_header(f) + except StopIteration: + break; + if not (id_ in element_types_names): + sys.stderr.write("mkvparse: Unknown element with id %x and size %d\n"%(id_, size)) + (resync_element_id, resync_element_size, resync_element_headersize) = resync(f) + if resync_element_id: + continue; + else: + break; + else: + id_ = resync_element_id + size=resync_element_size + hsize=resync_element_headersize + resync_element_id = None + resync_element_size = None + resync_element_headersize = None + + (type_, name) = element_types_names[id_] + (type_, name) = element_types_names[id_] + type_ = handler.begin_handling_ebml_element(id_, name, type_, hsize, size) + + if type_ == EET.MASTER: + tree = read_ebml_element_tree(f, size) + data = tree + + except Exception: + traceback.print_exc() + handler.before_handling_an_element() + (resync_element_id, resync_element_size, resync_element_headersize) = resync(f) + if resync_element_id: + continue; + else: + break; + + if name=="EBML" and type(data) == list: + d = dict(tree) + if 'EBMLReadVersion' in d: + if d['EBMLReadVersion'][1]>1: sys.stderr.write("mkvparse: Warning: EBMLReadVersion too big\n") + if 'DocTypeReadVersion' in d: + if d['DocTypeReadVersion'][1]>2: sys.stderr.write("mkvparse: Warning: DocTypeReadVersion too big\n") + dt = d['DocType'][1] + if dt != "matroska" and dt != "webm": + sys.stderr.write("mkvparse: Warning: EBML DocType is not \"matroska\" or \"webm\"") + elif name=="Info" and type(data) == list: + handler.segment_info = tree + handler.segment_info_available() + + d = dict(tree) + if "TimecodeScale" in d: + timecode_scale = d["TimecodeScale"][1] + elif name=="Tracks" and type(data) == list: + handler.tracks={} + for (ten, (_t, track)) in tree: + if ten != "TrackEntry": continue + d = dict(track) + n = d['TrackNumber'][1] + handler.tracks[n]=d + tt = d['TrackType'][1] + if tt==0x01: d['type']='video' + elif tt==0x02: d['type']='audio' + elif tt==0x03: d['type']='complex' + elif tt==0x10: d['type']='logo' + elif tt==0x11: d['type']='subtitle' + elif tt==0x12: d['type']='button' + elif tt==0x20: d['type']='control' + if 'TrackTimecodeScale' in d: + sys.stderr.write("mkvparse: Warning: TrackTimecodeScale is not supported\n") + if 'ContentEncodings' in d: + try: + compr = dict(d["ContentEncodings"][1][0][1][1][0][1][1]) + if compr["ContentCompAlgo"][1] == 3: + header_removal_headers_for_tracks[n] = compr["ContentCompSettings"][1] + else: + sys.stderr.write("mkvparse: Warning: compression other than " \ + "header removal is not supported\n") + except: + sys.stderr.write("mkvparse: Warning: unsuccessfully tried " \ + "to handle header removal compression\n") + handler.tracks_available() + # cluster contents: + elif name=="Timecode" and type_ == EET.UNSIGNED: + data=read_fixedlength_number(f, size, False) + current_cluster_timecode = data; + elif name=="SimpleBlock" and type_ == EET.BINARY: + pos = f.tell() + data=f.read(size) + handle_block(data, pos, handler, current_cluster_timecode, timecode_scale, None, header_removal_headers_for_tracks) + elif name=="BlockGroup" and type_ == EET.MASTER: + d2 = dict(tree) + duration=None + raise NotImplementedError + # if 'BlockDuration' in d2: + # duration = d2['BlockDuration'][1] + # duration = duration*0.000000001*timecode_scale + # if 'Block' in d2: + # handle_block(d2['Block'][1], None, handler, current_cluster_timecode, timecode_scale, duration, header_removal_headers_for_tracks) + else: + if type_!=EET.JUST_GO_ON and type_!=EET.MASTER: + data = read_simple_element(f, type_, size) + + handler.ebml_top_element(id_, name, type_, data); + + + +if __name__ == '__main__': + print("Run mkvuser.py for the example") diff --git a/tools/lib/pollable_queue.py b/tools/lib/pollable_queue.py new file mode 100644 index 0000000000..9ef2db62b2 --- /dev/null +++ b/tools/lib/pollable_queue.py @@ -0,0 +1,107 @@ +import os +import select +import fcntl +from itertools import count +from collections import deque + +Empty = OSError +Full = OSError +ExistentialError = OSError + +class PollableQueue(object): + """A Queue that you can poll(). + Only works with a single producer. + """ + def __init__(self, maxlen=None): + with open("/proc/sys/fs/pipe-max-size") as f: + max_maxlen = int(f.read().rstrip()) + + if maxlen is None: + maxlen = max_maxlen + else: + maxlen = min(maxlen, max_maxlen) + + self._maxlen = maxlen + self._q = deque() + self._get_fd, self._put_fd = os.pipe() + fcntl.fcntl(self._get_fd, fcntl.F_SETFL, os.O_NONBLOCK) + fcntl.fcntl(self._put_fd, fcntl.F_SETFL, os.O_NONBLOCK) + + fcntl.fcntl(self._get_fd, fcntl.F_SETLEASE + 7, self._maxlen) + fcntl.fcntl(self._put_fd, fcntl.F_SETLEASE + 7, self._maxlen) + + get_poller = select.epoll() + put_poller = select.epoll() + get_poller.register(self._get_fd, select.EPOLLIN) + put_poller.register(self._put_fd, select.EPOLLOUT) + + self._get_poll = get_poller.poll + self._put_poll = put_poller.poll + + + def get_fd(self): + return self._get_fd + + def put_fd(self): + return self._put_fd + + def put(self, item, block=True, timeout=None): + if block: + while self._put_poll(timeout if timeout is not None else -1): + try: + # TODO: This is broken for multiple push threads when the queue is full. + return self.put_nowait(item) + except OSError as e: + if e.errno != 11: + raise + + raise Full() + else: + return self.put_nowait(item) + + def put_nowait(self, item): + self._q.appendleft(item) + os.write(self._put_fd, b"\x00") + + def get(self, block=True, timeout=None): + if block: + while self._get_poll(timeout if timeout is not None else -1): + try: + return self.get_nowait() + except OSError as e: + if e.errno != 11: + raise + + raise Empty() + else: + return self.get_nowait() + + def get_nowait(self): + os.read(self._get_fd, 1) + return self._q.pop() + + def get_multiple(self, block=True, timeout=None): + if block: + if self._get_poll(timeout if timeout is not None else -1): + return self.get_multiple_nowait() + else: + raise Empty() + else: + return self.get_multiple_nowait() + + def get_multiple_nowait(self, max_messages=None): + num_read = len(os.read(self._get_fd, max_messages or self._maxlen)) + return [self._q.pop() for _ in range(num_read)] + + def empty(self): + return len(self._q) == 0 + + def full(self): + return len(self._q) >= self._maxlen + + def close(self): + os.close(self._get_fd) + os.close(self._put_fd) + + def __len__(self): + return len(self._q) \ No newline at end of file diff --git a/tools/lib/route.py b/tools/lib/route.py new file mode 100644 index 0000000000..359c1f22bd --- /dev/null +++ b/tools/lib/route.py @@ -0,0 +1,97 @@ +import os +import re +from collections import defaultdict + +SEGMENT_NAME_RE = r'[a-z0-9]{16}[|_][0-9]{4}-[0-9]{2}-[0-9]{2}--[0-9]{2}-[0-9]{2}-[0-9]{2}--[0-9]+' +EXPLORER_FILE_RE = r'^({})--([a-z]+\.[a-z0-9]+)$'.format(SEGMENT_NAME_RE) +OP_SEGMENT_DIR_RE = r'^({})$'.format(SEGMENT_NAME_RE) + +LOG_FILENAMES = ['rlog.bz2', 'raw_log.bz2', 'log2.gz', 'ilog.7z'] +CAMERA_FILENAMES = ['fcamera.hevc', 'video.hevc', 'acamera', 'icamera'] + +class Route(object): + def __init__(self, route_name, data_dir): + self.route_name = route_name.replace('_', '|') + self._segments = self._get_segments(data_dir) + + @property + def segments(self): + return self._segments + + def log_paths(self): + max_seg_number = self._segments[-1].canonical_name.segment_num + log_path_by_seg_num = {s.canonical_name.segment_num: s.log_path for s in self._segments} + return [log_path_by_seg_num.get(i, None) for i in range(max_seg_number+1)] + + def camera_paths(self): + max_seg_number = self._segments[-1].canonical_name.segment_num + camera_path_by_seg_num = {s.canonical_name.segment_num: s.camera_path for s in self._segments} + return [camera_path_by_seg_num.get(i, None) for i in range(max_seg_number+1)] + + def _get_segments(self, data_dir): + files = os.listdir(data_dir) + segment_files = defaultdict(list) + + for f in files: + fullpath = os.path.join(data_dir, f) + explorer_match = re.match(EXPLORER_FILE_RE, f) + op_match = re.match(OP_SEGMENT_DIR_RE, f) + + if explorer_match: + segment_name, fn = explorer_match.groups() + if segment_name.replace('_', '|').startswith(self.route_name): + segment_files[segment_name].append((fullpath, fn)) + elif op_match and os.path.isdir(fullpath): + segment_name, = op_match.groups() + if segment_name.startswith(self.route_name): + for seg_f in os.listdir(fullpath): + segment_files[segment_name].append((os.path.join(fullpath, seg_f), seg_f)) + elif f == self.route_name: + for seg_num in os.listdir(fullpath): + if not seg_num.isdigit(): + continue + + segment_name = '{}--{}'.format(self.route_name, seg_num) + for seg_f in os.listdir(os.path.join(fullpath, seg_num)): + segment_files[segment_name].append((os.path.join(fullpath, seg_num, seg_f), seg_f)) + + segments = [] + for segment, files in segment_files.items(): + try: + log_path = next(path for path, filename in files if filename in LOG_FILENAMES) + except StopIteration: + log_path = None + + try: + camera_path = next(path for path, filename in files if filename in CAMERA_FILENAMES) + except StopIteration: + camera_path = None + + segments.append(RouteSegment(segment, log_path, camera_path)) + + if len(segments) == 0: + raise ValueError('Could not find segments for route {} in data directory {}'.format(self.route_name, data_dir)) + return sorted(segments, key=lambda seg: seg.canonical_name.segment_num) + +class RouteSegment(object): + def __init__(self, name, log_path, camera_path): + self._name = RouteSegmentName(name) + self.log_path = log_path + self.camera_path = camera_path + + @property + def name(self): return str(self._name) + + @property + def canonical_name(self): return self._name + +class RouteSegmentName(object): + def __init__(self, name_str): + self._segment_name_str = name_str + self._route_name_str, num_str = self._segment_name_str.rsplit("--", 1) + self._num = int(num_str) + + @property + def segment_num(self): return self._num + + def __str__(self): return self._segment_name_str diff --git a/tools/lib/route_framereader.py b/tools/lib/route_framereader.py new file mode 100644 index 0000000000..47250383c5 --- /dev/null +++ b/tools/lib/route_framereader.py @@ -0,0 +1,86 @@ +"""RouteFrameReader indexes and reads frames across routes, by frameId or segment indices.""" +from tools.lib.framereader import FrameReader + +class _FrameReaderDict(dict): + def __init__(self, camera_paths, cache_paths, framereader_kwargs, *args, **kwargs): + super(_FrameReaderDict, self).__init__(*args, **kwargs) + + if cache_paths is None: + cache_paths = {} + if not isinstance(cache_paths, dict): + cache_paths = { k: v for k, v in enumerate(cache_paths) } + + self._camera_paths = camera_paths + self._cache_paths = cache_paths + self._framereader_kwargs = framereader_kwargs + + def __missing__(self, key): + if key < len(self._camera_paths) and self._camera_paths[key] is not None: + frame_reader = FrameReader(self._camera_paths[key], + self._cache_paths.get(key), **self._framereader_kwargs) + self[key] = frame_reader + return frame_reader + else: + raise KeyError("Segment index out of bounds: {}".format(key)) + + +class RouteFrameReader(object): + """Reads frames across routes and route segments by frameId.""" + def __init__(self, camera_paths, cache_paths, frame_id_lookup, **kwargs): + """Create a route framereader. + + Inputs: + TODO + + kwargs: Forwarded to the FrameReader function. If cache_prefix is included, that path + will also be used for frame position indices. + """ + self._first_camera_idx = next(i for i in range(len(camera_paths)) if camera_paths[i] is not None) + self._frame_readers = _FrameReaderDict(camera_paths, cache_paths, kwargs) + self._frame_id_lookup = frame_id_lookup + + @property + def w(self): + """Width of each frame in pixels.""" + return self._frame_readers[self._first_camera_idx].w + + @property + def h(self): + """Height of each frame in pixels.""" + return self._frame_readers[self._first_camera_idx].h + + def get(self, frame_id, **kwargs): + """Get a frame for a route based on frameId. + + Inputs: + frame_id: The frameId of the returned frame. + kwargs: Forwarded to BaseFrameReader.get. "count" is not implemented. + """ + segment_num, segment_id = self._frame_id_lookup.get(frame_id, (None, None)) + if segment_num is None or segment_num == -1 or segment_id == -1: + return None + else: + return self.get_from_segment(segment_num, segment_id, **kwargs) + + def get_from_segment(self, segment_num, segment_id, **kwargs): + """Get a frame from a specific segment with a specific index in that segment (segment_id). + + Inputs: + segment_num: The number of the segment. + segment_id: The index of the return frame within that segment. + kwargs: Forwarded to BaseFrameReader.get. "count" is not implemented. + """ + if "count" in kwargs: + raise NotImplementedError("count") + + return self._frame_readers[segment_num].get(segment_id, **kwargs)[0] + + + def close(self): + frs = self._frame_readers + self._frame_readers.clear() + for fr in frs: + fr.close() + + def __enter__(self): return self + def __exit__(self, type, value, traceback): self.close() diff --git a/tools/lib/tests/test_readers.py b/tools/lib/tests/test_readers.py new file mode 100755 index 0000000000..70e634909d --- /dev/null +++ b/tools/lib/tests/test_readers.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python +import unittest +import requests +import tempfile + +from collections import defaultdict +import numpy as np +from tools.lib.framereader import FrameReader +from tools.lib.logreader import LogReader + +class TestReaders(unittest.TestCase): + def test_logreader(self): + with tempfile.NamedTemporaryFile(suffix=".bz2") as fp: + r = requests.get("https://github.com/commaai/comma2k19/blob/master/Example_1/b0c9d2329ad1606b%7C2018-08-02--08-34-47/40/raw_log.bz2?raw=true") + fp.write(r.content) + fp.flush() + + lr = LogReader(fp.name) + hist = defaultdict(int) + for l in lr: + hist[l.which()] += 1 + + self.assertEqual(hist['carControl'], 6000) + self.assertEqual(hist['logMessage'], 6857) + + def test_framereader(self): + with tempfile.NamedTemporaryFile(suffix=".hevc") as fp: + r = requests.get("https://github.com/commaai/comma2k19/blob/master/Example_1/b0c9d2329ad1606b%7C2018-08-02--08-34-47/40/video.hevc?raw=true") + fp.write(r.content) + fp.flush() + + f = FrameReader(fp.name) + + self.assertEqual(f.frame_count, 1200) + self.assertEqual(f.w, 1164) + self.assertEqual(f.h, 874) + + + frame_first_30 = f.get(0, 30) + self.assertEqual(len(frame_first_30), 30) + + + print(frame_first_30[15]) + + print("frame_0") + frame_0 = f.get(0, 1) + frame_15 = f.get(15, 1) + + print(frame_15[0]) + + assert np.all(frame_first_30[0] == frame_0[0]) + assert np.all(frame_first_30[15] == frame_15[0]) + +if __name__ == "__main__": + unittest.main() + diff --git a/tools/lib/vidindex/.gitignore b/tools/lib/vidindex/.gitignore new file mode 100644 index 0000000000..a77a06e97d --- /dev/null +++ b/tools/lib/vidindex/.gitignore @@ -0,0 +1 @@ +vidindex diff --git a/tools/lib/vidindex/Makefile b/tools/lib/vidindex/Makefile new file mode 100644 index 0000000000..f6526db212 --- /dev/null +++ b/tools/lib/vidindex/Makefile @@ -0,0 +1,6 @@ +CC := gcc + +vidindex: bitstream.c bitstream.h vidindex.c + $(eval $@_TMP := $(shell mktemp)) + $(CC) -std=c99 bitstream.c vidindex.c -o $($@_TMP) + mv $($@_TMP) $@ diff --git a/tools/lib/vidindex/bitstream.c b/tools/lib/vidindex/bitstream.c new file mode 100644 index 0000000000..d174ffa8a7 --- /dev/null +++ b/tools/lib/vidindex/bitstream.c @@ -0,0 +1,118 @@ +#include +#include + +#include "bitstream.h" + +static const uint32_t BS_MASKS[33] = { + 0, 0x1L, 0x3L, 0x7L, 0xFL, 0x1FL, + 0x3FL, 0x7FL, 0xFFL, 0x1FFL, 0x3FFL, 0x7FFL, + 0xFFFL, 0x1FFFL, 0x3FFFL, 0x7FFFL, 0xFFFFL, 0x1FFFFL, + 0x3FFFFL, 0x7FFFFL, 0xFFFFFL, 0x1FFFFFL, 0x3FFFFFL, 0x7FFFFFL, + 0xFFFFFFL, 0x1FFFFFFL, 0x3FFFFFFL, 0x7FFFFFFL, 0xFFFFFFFL, 0x1FFFFFFFL, + 0x3FFFFFFFL, 0x7FFFFFFFL, 0xFFFFFFFFL}; + +void bs_init(struct bitstream* bs, const uint8_t* buffer, size_t input_size) { + bs->buffer_ptr = buffer; + bs->buffer_end = buffer + input_size; + bs->value = 0; + bs->pos = 0; + bs->shift = 8; + bs->size = input_size * 8; +} + +uint32_t bs_get(struct bitstream* bs, int n) { + if (n > 32) + return 0; + + bs->pos += n; + bs->shift += n; + while (bs->shift > 8) { + if (bs->buffer_ptr < bs->buffer_end) { + bs->value <<= 8; + bs->value |= *bs->buffer_ptr++; + bs->shift -= 8; + } else { + bs_seek(bs, bs->pos - n); + return 0; + // bs->value <<= 8; + // bs->shift -= 8; + } + } + return (bs->value >> (8 - bs->shift)) & BS_MASKS[n]; +} + +void bs_seek(struct bitstream* bs, size_t new_pos) { + bs->pos = (new_pos / 32) * 32; + bs->shift = 8; + bs->value = 0; + bs_get(bs, new_pos % 32); +} + +uint32_t bs_peek(struct bitstream* bs, int n) { + struct bitstream bak = *bs; + return bs_get(&bak, n); +} + +size_t bs_remain(struct bitstream* bs) { + return bs->size - bs->pos; +} + +int bs_eof(struct bitstream* bs) { + return bs_remain(bs) == 0; +} + +uint32_t bs_ue(struct bitstream* bs) { + static const uint8_t exp_golomb_bits[256] = { + 8, 7, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 4, 4, 4, 4, 3, 3, 3, 3, 3, 3, 3, 3, + 3, 3, 3, 3, 3, 3, 3, 3, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, + 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, + 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, + }; + uint32_t bits, read = 0; + int bits_left; + uint8_t coded; + int done = 0; + bits = 0; + // we want to read 8 bits at a time - if we don't have 8 bits, + // read what's left, and shift. The exp_golomb_bits calc remains the + // same. + while (!done) { + bits_left = bs_remain(bs); + if (bits_left < 8) { + read = bs_peek(bs, bits_left) << (8 - bits_left); + done = 1; + } else { + read = bs_peek(bs, 8); + if (read == 0) { + bs_get(bs, 8); + bits += 8; + } else { + done = 1; + } + } + } + coded = exp_golomb_bits[read]; + bs_get(bs, coded); + bits += coded; + + // printf("ue - bits %d\n", bits); + return bs_get(bs, bits + 1) - 1; +} + +int32_t bs_se(struct bitstream* bs) { + uint32_t ret; + ret = bs_ue(bs); + if ((ret & 0x1) == 0) { + ret >>= 1; + int32_t temp = 0 - ret; + return temp; + } + return (ret + 1) >> 1; +} diff --git a/tools/lib/vidindex/bitstream.h b/tools/lib/vidindex/bitstream.h new file mode 100644 index 0000000000..0f538a59ab --- /dev/null +++ b/tools/lib/vidindex/bitstream.h @@ -0,0 +1,26 @@ +#ifndef bitstream_H +#define bitstream_H + + +#include +#include + +struct bitstream { + const uint8_t *buffer_ptr; + const uint8_t *buffer_end; + uint64_t value; + uint32_t pos; + uint32_t shift; + size_t size; +}; + +void bs_init(struct bitstream *bs, const uint8_t *buffer, size_t input_size); +void bs_seek(struct bitstream *bs, size_t new_pos); +uint32_t bs_get(struct bitstream *bs, int n); +uint32_t bs_peek(struct bitstream *bs, int n); +size_t bs_remain(struct bitstream *bs); +int bs_eof(struct bitstream *bs); +uint32_t bs_ue(struct bitstream *bs); +int32_t bs_se(struct bitstream *bs); + +#endif diff --git a/tools/lib/vidindex/vidindex.c b/tools/lib/vidindex/vidindex.c new file mode 100644 index 0000000000..a8d53d947e --- /dev/null +++ b/tools/lib/vidindex/vidindex.c @@ -0,0 +1,307 @@ +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "bitstream.h" + +#define START_CODE 0x000001 + +static uint32_t read24be(const uint8_t* ptr) { + return (ptr[0] << 16) | (ptr[1] << 8) | ptr[2]; +} +static void write32le(FILE *of, uint32_t v) { + uint8_t va[4] = { + v & 0xff, (v >> 8) & 0xff, (v >> 16) & 0xff, (v >> 24) & 0xff + }; + fwrite(va, 1, sizeof(va), of); +} + +// Table 7-1 +enum hevc_nal_type { + HEVC_NAL_TYPE_TRAIL_N = 0, + HEVC_NAL_TYPE_TRAIL_R = 1, + HEVC_NAL_TYPE_TSA_N = 2, + HEVC_NAL_TYPE_TSA_R = 3, + HEVC_NAL_TYPE_STSA_N = 4, + HEVC_NAL_TYPE_STSA_R = 5, + HEVC_NAL_TYPE_RADL_N = 6, + HEVC_NAL_TYPE_RADL_R = 7, + HEVC_NAL_TYPE_RASL_N = 8, + HEVC_NAL_TYPE_RASL_R = 9, + HEVC_NAL_TYPE_BLA_W_LP = 16, + HEVC_NAL_TYPE_BLA_W_RADL = 17, + HEVC_NAL_TYPE_BLA_N_LP = 18, + HEVC_NAL_TYPE_IDR_W_RADL = 19, + HEVC_NAL_TYPE_IDR_N_LP = 20, + HEVC_NAL_TYPE_CRA_NUT = 21, + HEVC_NAL_TYPE_RSV_IRAP_VCL23 = 23, + HEVC_NAL_TYPE_VPS_NUT = 32, + HEVC_NAL_TYPE_SPS_NUT = 33, + HEVC_NAL_TYPE_PPS_NUT = 34, + HEVC_NAL_TYPE_AUD_NUT = 35, + HEVC_NAL_TYPE_EOS_NUT = 36, + HEVC_NAL_TYPE_EOB_NUT = 37, + HEVC_NAL_TYPE_FD_NUT = 38, + HEVC_NAL_TYPE_PREFIX_SEI_NUT = 39, + HEVC_NAL_TYPE_SUFFIX_SEI_NUT = 40, +}; + +// Table 7-7 +enum hevc_slice_type { + HEVC_SLICE_B = 0, + HEVC_SLICE_P = 1, + HEVC_SLICE_I = 2, +}; + +static void hevc_index(const uint8_t *data, size_t file_size, FILE *of_prefix, FILE *of_index) { + const uint8_t* ptr = data; + const uint8_t* ptr_end = data + file_size; + + assert(ptr[0] == 0); + ptr++; + assert(read24be(ptr) == START_CODE); + + // pps. ignore for now + uint32_t num_extra_slice_header_bits = 0; + uint32_t dependent_slice_segments_enabled_flag = 0; + + while (ptr < ptr_end) { + const uint8_t* next = ptr+1; + for (; next < ptr_end-4; next++) { + if (read24be(next) == START_CODE) break; + } + size_t nal_size = next - ptr; + if (nal_size < 6) { + break; + } + + { + struct bitstream bs = {0}; + bs_init(&bs, ptr, nal_size); + + uint32_t start_code = bs_get(&bs, 24); + assert(start_code == 0x000001); + + // nal_unit_header + uint32_t forbidden_zero_bit = bs_get(&bs, 1); + uint32_t nal_unit_type = bs_get(&bs, 6); + uint32_t nuh_layer_id = bs_get(&bs, 6); + uint32_t nuh_temporal_id_plus1 = bs_get(&bs, 3); + + // if (nal_unit_type != 1) printf("%3d -- %3d %10d %lu\n", nal_unit_type, frame_num, (uint32_t)(ptr-data), nal_size); + + switch (nal_unit_type) { + case HEVC_NAL_TYPE_VPS_NUT: + case HEVC_NAL_TYPE_SPS_NUT: + case HEVC_NAL_TYPE_PPS_NUT: + fwrite(ptr, 1, nal_size, of_prefix); + break; + case HEVC_NAL_TYPE_TRAIL_N: + case HEVC_NAL_TYPE_TRAIL_R: + case HEVC_NAL_TYPE_TSA_N: + case HEVC_NAL_TYPE_TSA_R: + case HEVC_NAL_TYPE_STSA_N: + case HEVC_NAL_TYPE_STSA_R: + case HEVC_NAL_TYPE_RADL_N: + case HEVC_NAL_TYPE_RADL_R: + case HEVC_NAL_TYPE_RASL_N: + case HEVC_NAL_TYPE_RASL_R: + case HEVC_NAL_TYPE_BLA_W_LP: + case HEVC_NAL_TYPE_BLA_W_RADL: + case HEVC_NAL_TYPE_BLA_N_LP: + case HEVC_NAL_TYPE_IDR_W_RADL: + case HEVC_NAL_TYPE_IDR_N_LP: + case HEVC_NAL_TYPE_CRA_NUT: { + // slice_segment_header + uint32_t first_slice_segment_in_pic_flag = bs_get(&bs, 1); + if (nal_unit_type >= HEVC_NAL_TYPE_BLA_W_LP && nal_unit_type <= HEVC_NAL_TYPE_RSV_IRAP_VCL23) { + uint32_t no_output_of_prior_pics_flag = bs_get(&bs, 1); + } + uint32_t slice_pic_parameter_set_id = bs_get(&bs, 1); + if (!first_slice_segment_in_pic_flag) { + // ... + break; + } + + if (!dependent_slice_segments_enabled_flag) { + for (int i=0; i 4); + + const uint8_t* data = (const uint8_t*)mmap(NULL, file_size, PROT_READ, MAP_PRIVATE, fd, 0); + assert(data != MAP_FAILED); + + if (strcmp(file_type, "hevc") == 0) { + hevc_index(data, file_size, of_prefix, of_index); + } else if (strcmp(file_type, "h264") == 0) { + h264_index(data, file_size, of_prefix, of_index); + } else { + assert(false); + } + + munmap((void*)data, file_size); + close(fd); + + return 0; +} diff --git a/tools/livedm/helpers.py b/tools/livedm/helpers.py new file mode 100644 index 0000000000..47a79a67cb --- /dev/null +++ b/tools/livedm/helpers.py @@ -0,0 +1,30 @@ +import numpy as np +import cv2 + +def rot_matrix(roll, pitch, yaw): + cr, sr = np.cos(roll), np.sin(roll) + cp, sp = np.cos(pitch), np.sin(pitch) + cy, sy = np.cos(yaw), np.sin(yaw) + rr = np.array([[1,0,0],[0, cr,-sr],[0, sr, cr]]) + rp = np.array([[cp,0,sp],[0, 1,0],[-sp, 0, cp]]) + ry = np.array([[cy,-sy,0],[sy, cy,0],[0, 0, 1]]) + return ry.dot(rp.dot(rr)) + +def draw_pose(img, pose, loc, W=160, H=320, xyoffset=(0,0), faceprob=0): + rcmat = np.zeros((3,4)) + rcmat[:,:3] = rot_matrix(*pose[0:3]) * 0.5 + rcmat[0,3] = (loc[0]+0.5) * W + rcmat[1,3] = (loc[1]+0.5) * H + rcmat[2,3] = 1.0 + # draw nose + p1 = np.dot(rcmat, [0,0,0,1])[0:2] + p2 = np.dot(rcmat, [0,0,100,1])[0:2] + tr = tuple([int(round(x + xyoffset[i])) for i,x in enumerate(p1)]) + pr = tuple([int(round(x + xyoffset[i])) for i,x in enumerate(p2)]) + if faceprob > 0.4: + color = (255,255,0) + cv2.line(img, tr, pr, color=(255,255,0), thickness=3) + else: + color = (64,64,64) + cv2.circle(img, tr, 7, color=color) + \ No newline at end of file diff --git a/tools/livedm/livedm.py b/tools/livedm/livedm.py new file mode 100644 index 0000000000..5125f0e2c7 --- /dev/null +++ b/tools/livedm/livedm.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +import os +import argparse +import pygame +import numpy as np +import cv2 + +from cereal import log +import cereal.messaging as messaging + +from helpers import draw_pose + +if __name__ == "__main__": + + os.environ["ZMQ"] = "1" + + parser = argparse.ArgumentParser(description='Sniff a communcation socket') + parser.add_argument('--addr', default='192.168.5.11') + args = parser.parse_args() + + messaging.context = messaging.Context() + + poller = messaging.Poller() + + m = 'driverMonitoring' + sock = messaging.sub_sock(m, poller, addr=args.addr) + + pygame.init() + pygame.display.set_caption('livedm') + screen = pygame.display.set_mode((320,640), pygame.DOUBLEBUF) + camera_surface = pygame.surface.Surface((160,320), 0, 24).convert() + + while 1: + polld = poller.poll(1000) + for sock in polld: + msg = sock.receive() + evt = log.Event.from_bytes(msg) + + faceProb = np.array(evt.driverMonitoring.faceProb) + faceOrientation = np.array(evt.driverMonitoring.faceOrientation) + facePosition = np.array(evt.driverMonitoring.facePosition) + + print(faceProb) + # print(faceOrientation) + # print(facePosition) + faceOrientation[1] *= -1 + facePosition[0] *= -1 + + img = np.zeros((320,160,3)) + if faceProb > 0.4: + cv2.putText(img, 'you', (int(facePosition[0]*160+40), int(facePosition[1]*320+110)), cv2.FONT_ITALIC, 0.5, (255,255,0)) + cv2.rectangle(img, (int(facePosition[0]*160+40), int(facePosition[1]*320+120)),\ + (int(facePosition[0]*160+120), int(facePosition[1]*320+200)), (255,255,0), 1) + + not_blink = evt.driverMonitoring.leftBlinkProb + evt.driverMonitoring.rightBlinkProb < 1 + + if evt.driverMonitoring.leftEyeProb > 0.6: + cv2.line(img, (int(facePosition[0]*160+95), int(facePosition[1]*320+140)),\ + (int(facePosition[0]*160+105), int(facePosition[1]*320+140)), (255,255,0), 2) + if not_blink: + cv2.line(img, (int(facePosition[0]*160+99), int(facePosition[1]*320+143)),\ + (int(facePosition[0]*160+101), int(facePosition[1]*320+143)), (255,255,0), 2) + + if evt.driverMonitoring.rightEyeProb > 0.6: + cv2.line(img, (int(facePosition[0]*160+55), int(facePosition[1]*320+140)),\ + (int(facePosition[0]*160+65), int(facePosition[1]*320+140)), (255,255,0), 2) + if not_blink: + cv2.line(img, (int(facePosition[0]*160+59), int(facePosition[1]*320+143)),\ + (int(facePosition[0]*160+61), int(facePosition[1]*320+143)), (255,255,0), 2) + + else: + cv2.putText(img, 'you not found', (int(facePosition[0]*160+40), int(facePosition[1]*320+110)), cv2.FONT_ITALIC, 0.5, (64,64,64)) + draw_pose(img, faceOrientation, facePosition, + W = 160, H = 320, xyoffset = (0, 0), faceprob=faceProb) + + pygame.surfarray.blit_array(camera_surface, img.swapaxes(0,1)) + camera_surface_2x = pygame.transform.scale2x(camera_surface) + screen.blit(camera_surface_2x, (0, 0)) + pygame.display.flip() diff --git a/tools/misc/save_ubloxraw_stream.py b/tools/misc/save_ubloxraw_stream.py new file mode 100644 index 0000000000..518c4ecaf6 --- /dev/null +++ b/tools/misc/save_ubloxraw_stream.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python +import argparse +import os +import sys +from common.basedir import BASEDIR +from tools.lib.logreader import MultiLogIterator +from tools.lib.route import Route + +os.environ['BASEDIR'] = BASEDIR + + +def get_arg_parser(): + parser = argparse.ArgumentParser( + description="Unlogging and save to file", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + parser.add_argument("data_dir", nargs='?', + help="Path to directory in which log and camera files are located.") + parser.add_argument("route_name", type=(lambda x: x.replace("#", "|")), nargs="?", + help="The route whose messages will be published.") + parser.add_argument("--out_path", nargs='?', default='/data/ubloxRaw.stream', + help="Output pickle file path") + return parser + + +def main(argv): + args = get_arg_parser().parse_args(sys.argv[1:]) + if not args.data_dir: + print('Data directory invalid.') + return + + if not args.route_name: + # Extract route name from path + args.route_name = os.path.basename(args.data_dir) + args.data_dir = os.path.dirname(args.data_dir) + + route = Route(args.route_name, args.data_dir) + lr = MultiLogIterator(route.log_paths(), wraparound=False) + + with open(args.out_path, 'wb') as f: + try: + done = False + i = 0 + while not done: + msg = next(lr) + if not msg: + break + smsg = msg.as_builder() + typ = smsg.which() + if typ == 'ubloxRaw': + f.write(smsg.to_bytes()) + i += 1 + except StopIteration: + print('All done') + print('Writed {} msgs'.format(i)) + + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/tools/nui/.gitignore b/tools/nui/.gitignore new file mode 100644 index 0000000000..634d1cc12c --- /dev/null +++ b/tools/nui/.gitignore @@ -0,0 +1,8 @@ +Makefile +.*.swp +*.o +nui +moc_* +.qmake.stash +nui.app/* + diff --git a/tools/nui/FileReader.cpp b/tools/nui/FileReader.cpp new file mode 100644 index 0000000000..7c207dff83 --- /dev/null +++ b/tools/nui/FileReader.cpp @@ -0,0 +1,138 @@ +#include "FileReader.hpp" +#include "FrameReader.hpp" + +#include + +FileReader::FileReader(const QString& file_) : file(file_) { +} + +void FileReader::process() { + timer.start(); + // TODO: Support reading files from the API + startRequest(QUrl("http://data.comma.life/"+file)); +} + +void FileReader::startRequest(const QUrl &url) { + qnam = new QNetworkAccessManager; + reply = qnam->get(QNetworkRequest(url)); + connect(reply, &QNetworkReply::finished, this, &FileReader::httpFinished); + connect(reply, &QIODevice::readyRead, this, &FileReader::readyRead); + qDebug() << "requesting" << url; +} + +void FileReader::httpFinished() { + if (reply->error()) { + qWarning() << reply->errorString(); + } + + const QVariant redirectionTarget = reply->attribute(QNetworkRequest::RedirectionTargetAttribute); + if (!redirectionTarget.isNull()) { + const QUrl redirectedUrl = redirectionTarget.toUrl(); + //qDebug() << "redirected to" << redirectedUrl; + startRequest(redirectedUrl); + } else { + qDebug() << "done in" << timer.elapsed() << "ms"; + done(); + } +} + +void FileReader::readyRead() { + QByteArray dat = reply->readAll(); + printf("got http ready read: %d\n", dat.size()); +} + +FileReader::~FileReader() { + +} + +LogReader::LogReader(const QString& file, Events *events_, QReadWriteLock* events_lock_, QMap > *eidx_) : + FileReader(file), events(events_), events_lock(events_lock_), eidx(eidx_) { + bStream.next_in = NULL; + bStream.avail_in = 0; + bStream.bzalloc = NULL; + bStream.bzfree = NULL; + bStream.opaque = NULL; + + int ret = BZ2_bzDecompressInit(&bStream, 0, 0); + if (ret != BZ_OK) qWarning() << "bz2 init failed"; + + // start with 64MB buffer + raw.resize(1024*1024*64); + + // auto increment? + bStream.next_out = raw.data(); + bStream.avail_out = raw.size(); + + // parsed no events yet + event_offset = 0; + + parser = new std::thread([&]() { + while (1) { + mergeEvents(cdled.get()); + } + }); +} + +void LogReader::mergeEvents(int dled) { + auto amsg = kj::arrayPtr((const capnp::word*)(raw.data() + event_offset), (dled-event_offset)/sizeof(capnp::word)); + Events events_local; + QMap > eidx_local; + + while (amsg.size() > 0) { + try { + capnp::FlatArrayMessageReader cmsg = capnp::FlatArrayMessageReader(amsg); + + // this needed? it is + capnp::FlatArrayMessageReader *tmsg = + new capnp::FlatArrayMessageReader(kj::arrayPtr(amsg.begin(), cmsg.getEnd())); + + amsg = kj::arrayPtr(cmsg.getEnd(), amsg.end()); + + cereal::Event::Reader event = tmsg->getRoot(); + events_local.insert(event.getLogMonoTime(), event); + + // hack + // TODO: rewrite with callback + if (event.which() == cereal::Event::ENCODE_IDX) { + auto ee = event.getEncodeIdx(); + eidx_local.insert(ee.getFrameId(), qMakePair(ee.getSegmentNum(), ee.getSegmentId())); + } + + // increment + event_offset = (char*)cmsg.getEnd() - raw.data(); + } catch (const kj::Exception& e) { + // partial messages trigger this + //qDebug() << e.getDescription().cStr(); + break; + } + } + + // merge in events + // TODO: add lock + events_lock->lockForWrite(); + *events += events_local; + eidx->unite(eidx_local); + events_lock->unlock(); + + printf("parsed %d into %d events with offset %d\n", dled, events->size(), event_offset); +} + +void LogReader::readyRead() { + QByteArray dat = reply->readAll(); + + bStream.next_in = dat.data(); + bStream.avail_in = dat.size(); + + while (bStream.avail_in > 0) { + int ret = BZ2_bzDecompress(&bStream); + if (ret != BZ_OK && ret != BZ_STREAM_END) { + qWarning() << "bz2 decompress failed"; + break; + } + qDebug() << "got" << dat.size() << "with" << bStream.avail_out << "size" << raw.size(); + } + + int dled = raw.size() - bStream.avail_out; + cdled.put(dled); +} + diff --git a/tools/nui/FileReader.hpp b/tools/nui/FileReader.hpp new file mode 100644 index 0000000000..6d53bdee8e --- /dev/null +++ b/tools/nui/FileReader.hpp @@ -0,0 +1,68 @@ +#ifndef FILEREADER_HPP +#define FILEREADER_HPP + +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include "cereal/gen/cpp/log.capnp.h" + +#include +#include "channel.hpp" + +class FileReader : public QObject { +Q_OBJECT +public: + FileReader(const QString& file_); + void startRequest(const QUrl &url); + ~FileReader(); + virtual void readyRead(); + void httpFinished(); + virtual void done() {}; +public slots: + void process(); +protected: + QNetworkReply *reply; +private: + QNetworkAccessManager *qnam; + QElapsedTimer timer; + QString file; +}; + +typedef QMultiMap Events; + +class LogReader : public FileReader { +Q_OBJECT +public: + LogReader(const QString& file, Events *, QReadWriteLock* events_lock_, QMap > *eidx_); + void readyRead(); + void done() { is_done = true; }; + bool is_done = false; +private: + bz_stream bStream; + + // backing store + QByteArray raw; + + std::thread *parser; + int event_offset; + channel cdled; + + // global + void mergeEvents(int dled); + Events *events; + QReadWriteLock* events_lock; + QMap > *eidx; +}; + +#endif + diff --git a/tools/nui/README b/tools/nui/README new file mode 100644 index 0000000000..3033bcadc0 --- /dev/null +++ b/tools/nui/README @@ -0,0 +1,9 @@ +== Ubuntu == + +sudo apt-get install capnproto libyaml-cpp-dev qt5-default + +== Mac == + +brew install qt5 ffmpeg capnp yaml-cpp zmq +brew link qt5 --force + diff --git a/tools/nui/Unlogger.cpp b/tools/nui/Unlogger.cpp new file mode 100644 index 0000000000..843ef51e3f --- /dev/null +++ b/tools/nui/Unlogger.cpp @@ -0,0 +1,182 @@ +#include +#include +#include +#include +#include + +// include the dynamic struct +#include "cereal/gen/cpp/car.capnp.c++" +#include "cereal/gen/cpp/log.capnp.c++" + +#include "Unlogger.hpp" + +#include +#include + +static inline uint64_t nanos_since_boot() { + struct timespec t; + clock_gettime(CLOCK_BOOTTIME, &t); + return t.tv_sec * 1000000000ULL + t.tv_nsec; +} + + +Unlogger::Unlogger(Events *events_, QReadWriteLock* events_lock_, QMap *frs_, int seek) + : events(events_), events_lock(events_lock_), frs(frs_) { + ctx = Context::create(); + YAML::Node service_list = YAML::LoadFile("../../cereal/service_list.yaml"); + + seek_request = seek*1e9; + + QStringList block = QString(getenv("BLOCK")).split(","); + qDebug() << "blocklist" << block; + + QStringList allow = QString(getenv("ALLOW")).split(","); + qDebug() << "allowlist" << allow; + + for (const auto& it : service_list) { + auto name = it.first.as(); + + if (allow[0].size() > 0 && !allow.contains(name.c_str())) { + qDebug() << "not allowing" << name.c_str(); + continue; + } + + if (block.contains(name.c_str())) { + qDebug() << "blocking" << name.c_str(); + continue; + } + + PubSocket *sock = PubSocket::create(ctx, name); + if (sock == NULL) { + qDebug() << "FAILED" << name.c_str(); + continue; + } + + qDebug() << name.c_str(); + + for (auto field: capnp::Schema::from().getFields()) { + std::string tname = field.getProto().getName(); + + if (tname == name) { + // TODO: I couldn't figure out how to get the which, only the index, hence this hack + int type = field.getIndex(); + if (type > 67) type--; // valid + type--; // logMonoTime + + //qDebug() << "here" << tname.c_str() << type << cereal::Event::CONTROLS_STATE; + socks.insert(type, sock); + } + } + } +} + +void Unlogger::process() { + qDebug() << "hello from unlogger thread"; + while (events->size() == 0) { + qDebug() << "waiting for events"; + QThread::sleep(1); + } + qDebug() << "got events"; + + // TODO: hack + if (seek_request != 0) { + seek_request += events->begin().key(); + while (events->lowerBound(seek_request) == events->end()) { + qDebug() << "waiting for desired time"; + QThread::sleep(1); + } + } + + QElapsedTimer timer; + timer.start(); + + uint64_t last_elapsed = 0; + + // loops + while (1) { + uint64_t t0 = (events->begin()+1).key(); + uint64_t t0r = timer.nsecsElapsed(); + qDebug() << "unlogging at" << t0; + + auto eit = events->lowerBound(t0); + while (eit != events->end()) { + while (paused) { + QThread::usleep(1000); + t0 = eit->getLogMonoTime(); + t0r = timer.nsecsElapsed(); + } + + if (seek_request != 0) { + t0 = seek_request; + qDebug() << "seeking to" << t0; + t0r = timer.nsecsElapsed(); + eit = events->lowerBound(t0); + seek_request = 0; + if (eit == events->end()) { + qWarning() << "seek off end"; + break; + } + } + + if (abs(((long long)tc-(long long)last_elapsed)) > 50e6) { + //qDebug() << "elapsed"; + emit elapsed(); + last_elapsed = tc; + } + + auto e = *eit; + auto type = e.which(); + uint64_t tm = e.getLogMonoTime(); + auto it = socks.find(type); + tc = tm; + if (it != socks.end()) { + long etime = tm-t0; + long rtime = timer.nsecsElapsed() - t0r; + long us_behind = ((etime-rtime)*1e-3)+0.5; + if (us_behind > 0) { + if (us_behind > 1e6) { + qWarning() << "OVER ONE SECOND BEHIND, HACKING" << us_behind; + us_behind = 0; + t0 = tm; + t0r = timer.nsecsElapsed(); + } + QThread::usleep(us_behind); + //qDebug() << "sleeping" << us_behind << etime << timer.nsecsElapsed(); + } + + capnp::MallocMessageBuilder msg; + msg.setRoot(e); + + auto ee = msg.getRoot(); + ee.setLogMonoTime(nanos_since_boot()); + + if (e.which() == cereal::Event::FRAME) { + auto fr = msg.getRoot().getFrame(); + + // TODO: better way? + auto it = eidx.find(fr.getFrameId()); + if (it != eidx.end()) { + auto pp = *it; + //qDebug() << fr.getFrameId() << pp; + + if (frs->find(pp.first) != frs->end()) { + auto frm = (*frs)[pp.first]; + auto data = frm->get(pp.second); + if (data != NULL) { + fr.setImage(kj::arrayPtr(data, frm->getRGBSize())); + } + } + } + } + + auto words = capnp::messageToFlatArray(msg); + auto bytes = words.asBytes(); + + // TODO: Can PubSocket take a const char? + (*it)->send((char*)bytes.begin(), bytes.size()); + } + ++eit; + } + } +} + diff --git a/tools/nui/Unlogger.hpp b/tools/nui/Unlogger.hpp new file mode 100644 index 0000000000..577632268e --- /dev/null +++ b/tools/nui/Unlogger.hpp @@ -0,0 +1,36 @@ +#ifndef UNLOGGER_HPP +#define UNLOGGER_HPP + +#include +#include +#include "messaging.hpp" +#include "FileReader.hpp" +#include "FrameReader.hpp" + +class Unlogger : public QObject { +Q_OBJECT + public: + Unlogger(Events *events_, QReadWriteLock* events_lock_, QMap *frs_, int seek); + uint64_t getCurrentTime() { return tc; } + void setSeekRequest(uint64_t seek_request_) { seek_request = seek_request_; } + void setPause(bool pause) { paused = pause; } + void togglePause() { paused = !paused; } + QMap > eidx; + public slots: + void process(); + signals: + void elapsed(); + void finished(); + private: + Events *events; + QReadWriteLock *events_lock; + QMap *frs; + QMap socks; + Context *ctx; + uint64_t tc = 0; + uint64_t seek_request = 0; + bool paused = false; +}; + +#endif + diff --git a/tools/nui/build.sh b/tools/nui/build.sh new file mode 100755 index 0000000000..d859d6f1d0 --- /dev/null +++ b/tools/nui/build.sh @@ -0,0 +1,4 @@ +#!/bin/bash +qmake +make -j + diff --git a/tools/nui/main.cpp b/tools/nui/main.cpp new file mode 100644 index 0000000000..4bae08ed46 --- /dev/null +++ b/tools/nui/main.cpp @@ -0,0 +1,216 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "FileReader.hpp" +#include "Unlogger.hpp" +#include "FrameReader.hpp" + +class Window : public QWidget { + public: + Window(QString route_, int seek); + bool addSegment(int i); + protected: + void keyPressEvent(QKeyEvent *event) override; + void mousePressEvent(QMouseEvent *event) override; + void paintEvent(QPaintEvent *event) override; + uint64_t ct; + Unlogger *unlogger; + private: + int timeToPixel(uint64_t ns); + uint64_t pixelToTime(int px); + QString route; + + QReadWriteLock events_lock; + Events events; + int last_event_size = 0; + + QMap lrs; + QMap frs; + + // cache the bar + QPixmap *px = NULL; + int seg_add = 0; + + QLineEdit *timeLE; +}; + +Window::Window(QString route_, int seek) : route(route_) { + timeLE = new QLineEdit(this); + timeLE->setPlaceholderText("Placeholder Text"); + timeLE->move(50, 650); + + QThread* thread = new QThread; + unlogger = new Unlogger(&events, &events_lock, &frs, seek); + unlogger->moveToThread(thread); + connect(thread, SIGNAL (started()), unlogger, SLOT (process())); + connect(unlogger, SIGNAL (elapsed()), this, SLOT (update())); + thread->start(); + + this->setFocusPolicy(Qt::StrongFocus); + + // add the first segment + addSegment(seek/60); +} + +bool Window::addSegment(int i) { + if (lrs.find(i) == lrs.end()) { + QString fn = QString("%1/%2/rlog.bz2").arg(route).arg(i); + + QThread* thread = new QThread; + lrs.insert(i, new LogReader(fn, &events, &events_lock, &unlogger->eidx)); + lrs[i]->moveToThread(thread); + connect(thread, SIGNAL (started()), lrs[i], SLOT (process())); + thread->start(); + //connect(lrs[i], SIGNAL (finished()), this, SLOT (update())); + + QString frn = QString("%1/%2/fcamera.hevc").arg(route).arg(i); + frs.insert(i, new FrameReader(qPrintable(frn))); + return true; + } + return false; +} + +#define PIXELS_PER_SEC 0.5 + +int Window::timeToPixel(uint64_t ns) { + // TODO: make this dynamic + return int(ns*1e-9*PIXELS_PER_SEC+0.5); +} + +uint64_t Window::pixelToTime(int px) { + // TODO: make this dynamic + //printf("%d\n", px); + return ((px+0.5)/PIXELS_PER_SEC) * 1e9; +} + +void Window::keyPressEvent(QKeyEvent *event) { + printf("keypress: %x\n", event->key()); + if (event->key() == Qt::Key_Space) unlogger->togglePause(); +} + +void Window::mousePressEvent(QMouseEvent *event) { + //printf("mouse event\n"); + if (event->button() == Qt::LeftButton) { + uint64_t t0 = events.begin().key(); + uint64_t tt = pixelToTime(event->x()); + int seg = int((tt*1e-9)/60); + printf("segment %d\n", seg); + addSegment(seg); + + //printf("seek to %lu\n", t0+tt); + unlogger->setSeekRequest(t0+tt); + } + this->update(); +} + +void Window::paintEvent(QPaintEvent *event) { + if (events.size() == 0) return; + + QElapsedTimer timer; + timer.start(); + + uint64_t t0 = events.begin().key(); + uint64_t t1 = (events.end()-1).key(); + + //p.drawRect(0, 0, 600, 100); + + // TODO: we really don't have to redraw this every time, only on updates to events + int this_event_size = events.size(); + if (last_event_size != this_event_size) { + if (px != NULL) delete px; + px = new QPixmap(1920, 600); + px->fill(QColor(0xd8, 0xd8, 0xd8)); + + QPainter tt(px); + tt.setBrush(Qt::cyan); + + int lt = -1; + int lvv = 0; + for (auto e : events) { + auto type = e.which(); + //printf("%lld %d\n", e.getLogMonoTime()-t0, type); + if (type == cereal::Event::CONTROLS_STATE) { + auto controlsState = e.getControlsState(); + uint64_t t = (e.getLogMonoTime()-t0); + float vEgo = controlsState.getVEgo(); + int enabled = controlsState.getState() == cereal::ControlsState::OpenpilotState::ENABLED; + int rt = timeToPixel(t); // 250 ms per pixel + if (rt != lt) { + int vv = vEgo*8.0; + if (lt != -1) { + tt.setPen(Qt::red); + tt.drawLine(lt, 300-lvv, rt, 300-vv); + + if (enabled) { + tt.setPen(Qt::green); + } else { + tt.setPen(Qt::blue); + } + + tt.drawLine(rt, 300, rt, 600); + } + lt = rt; + lvv = vv; + } + } + } + tt.end(); + last_event_size = this_event_size; + if (lrs.find(seg_add) != lrs.end() && lrs[seg_add]->is_done) { + while (!addSegment(++seg_add)); + } + } + + QPainter p(this); + if (px != NULL) p.drawPixmap(0, 0, 1920, 600, *px); + + p.setBrush(Qt::cyan); + + uint64_t ct = unlogger->getCurrentTime(); + if (ct != 0) { + addSegment((((ct-t0)*1e-9)/60)+1); + int rrt = timeToPixel(ct-t0); + p.drawRect(rrt-1, 0, 2, 600); + + timeLE->setText(QString("%1").arg((ct-t0)*1e-9, '8', 'f', 2)); + } + + p.end(); + + if (timer.elapsed() > 50) { + qDebug() << "paint in" << timer.elapsed() << "ms"; + } +} + +int main(int argc, char *argv[]) { + QApplication app(argc, argv); + + QString route(argv[1]); + int seek = QString(argv[2]).toInt(); + printf("seek: %d\n", seek); + route = route.replace("|", "/"); + if (route == "") { + printf("usage %s: \n", argv[0]); + exit(0); + //route = "3a5d6ac1c23e5536/2019-10-29--10-06-58"; + //route = "0006c839f32a6f99/2019-02-18--06-21-29"; + //route = "02ec6bea180a4d36/2019-10-25--10-18-09"; + } + + Window window(route, seek); + window.resize(1920, 800); + window.setWindowTitle("nui unlogger"); + window.show(); + + return app.exec(); +} + diff --git a/tools/nui/nui.pro b/tools/nui/nui.pro new file mode 100644 index 0000000000..2a3c91439f --- /dev/null +++ b/tools/nui/nui.pro @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:69b914687867cabc35e1ba333479fd2f1d976d813e6d97fad48b94e5b470b57b +size 1105 diff --git a/tools/nui/test/.gitignore b/tools/nui/test/.gitignore new file mode 100644 index 0000000000..9daeafb986 --- /dev/null +++ b/tools/nui/test/.gitignore @@ -0,0 +1 @@ +test diff --git a/tools/nui/test/TestFrameReader.cpp b/tools/nui/test/TestFrameReader.cpp new file mode 100644 index 0000000000..944dbd4416 --- /dev/null +++ b/tools/nui/test/TestFrameReader.cpp @@ -0,0 +1,14 @@ +#include "FrameReader.hpp" +#include "TestFrameReader.hpp" + +void TestFrameReader::frameread() { + QElapsedTimer t; + t.start(); + FrameReader fr("3a5d6ac1c23e5536/2019-10-29--10-06-58/2/fcamera.hevc"); + fr.get(2); + //QThread::sleep(10); + qDebug() << t.nsecsElapsed()*1e-9 << "seconds"; +} + +QTEST_MAIN(TestFrameReader) + diff --git a/tools/nui/test/TestFrameReader.hpp b/tools/nui/test/TestFrameReader.hpp new file mode 100644 index 0000000000..59c07a3ebd --- /dev/null +++ b/tools/nui/test/TestFrameReader.hpp @@ -0,0 +1,8 @@ +#include + +class TestFrameReader : public QObject { +Q_OBJECT +private slots: + void frameread(); +}; + diff --git a/tools/nui/test/test.pro b/tools/nui/test/test.pro new file mode 100644 index 0000000000..72bbe16230 --- /dev/null +++ b/tools/nui/test/test.pro @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:5f7da1fd147df1bf851d759749e2b160bcfc6e4f7557a5b41653f85e4a0c029a +size 447 diff --git a/tools/replay/__init__.py b/tools/replay/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tools/replay/camera.py b/tools/replay/camera.py new file mode 100755 index 0000000000..a48c006407 --- /dev/null +++ b/tools/replay/camera.py @@ -0,0 +1,112 @@ +#!/usr/bin/env python +import os + +from common.basedir import BASEDIR +os.environ['BASEDIR'] = BASEDIR +SCALE = 3 + +import argparse +import zmq +import pygame +import numpy as np +import cv2 +import sys +import traceback +from collections import namedtuple +from cereal import car +from common.params import Params +from tools.lib.lazy_property import lazy_property +from cereal.messaging import sub_sock, recv_one_or_none, recv_one +from cereal.services import service_list + +_BB_OFFSET = 0, 0 +_BB_TO_FULL_FRAME = np.asarray([[1., 0., _BB_OFFSET[0]], [0., 1., _BB_OFFSET[1]], + [0., 0., 1.]]) +_FULL_FRAME_TO_BB = np.linalg.inv(_BB_TO_FULL_FRAME) +_FULL_FRAME_SIZE = 1164, 874 + + + +def pygame_modules_have_loaded(): + return pygame.display.get_init() and pygame.font.get_init() + + +def ui_thread(addr, frame_address): + context = zmq.Context.instance() + + pygame.init() + pygame.font.init() + assert pygame_modules_have_loaded() + + size = (_FULL_FRAME_SIZE[0] * SCALE, _FULL_FRAME_SIZE[1] * SCALE) + pygame.display.set_caption("comma one debug UI") + screen = pygame.display.set_mode(size, pygame.DOUBLEBUF) + + camera_surface = pygame.surface.Surface((_FULL_FRAME_SIZE[0] * SCALE, _FULL_FRAME_SIZE[1] * SCALE), 0, 24).convert() + + frame = context.socket(zmq.SUB) + frame.connect(frame_address or "tcp://%s:%d" % (addr, service_list['frame'].port)) + frame.setsockopt(zmq.SUBSCRIBE, "") + + img = np.zeros((_FULL_FRAME_SIZE[1], _FULL_FRAME_SIZE[0], 3), dtype='uint8') + imgff = np.zeros((_FULL_FRAME_SIZE[1], _FULL_FRAME_SIZE[0], 3), dtype=np.uint8) + + while 1: + list(pygame.event.get()) + screen.fill((64, 64, 64)) + + # ***** frame ***** + fpkt = recv_one(frame) + yuv_img = fpkt.frame.image + + if fpkt.frame.transform: + yuv_transform = np.array(fpkt.frame.transform).reshape(3, 3) + else: + # assume frame is flipped + yuv_transform = np.array([[-1.0, 0.0, _FULL_FRAME_SIZE[0] - 1], + [0.0, -1.0, _FULL_FRAME_SIZE[1] - 1], [0.0, 0.0, 1.0]]) + + if yuv_img and len(yuv_img) == _FULL_FRAME_SIZE[0] * _FULL_FRAME_SIZE[1] * 3 // 2: + yuv_np = np.frombuffer( + yuv_img, dtype=np.uint8).reshape(_FULL_FRAME_SIZE[1] * 3 // 2, -1) + cv2.cvtColor(yuv_np, cv2.COLOR_YUV2RGB_I420, dst=imgff) + cv2.warpAffine( + imgff, + np.dot(yuv_transform, _BB_TO_FULL_FRAME)[:2], (img.shape[1], img.shape[0]), + dst=img, + flags=cv2.WARP_INVERSE_MAP) + else: + img.fill(0) + + height, width = img.shape[:2] + img_resized = cv2.resize( + img, (SCALE * width, SCALE * height), interpolation=cv2.INTER_CUBIC) + # *** blits *** + pygame.surfarray.blit_array(camera_surface, img_resized.swapaxes(0, 1)) + screen.blit(camera_surface, (0, 0)) + + # this takes time...vsync or something + pygame.display.flip() + + +def get_arg_parser(): + parser = argparse.ArgumentParser( + description="Show replay data in a UI.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + parser.add_argument( + "ip_address", + nargs="?", + default="127.0.0.1", + help="The ip address on which to receive zmq messages.") + + parser.add_argument( + "--frame-address", + default=None, + help="The ip address on which to receive zmq messages.") + return parser + + +if __name__ == "__main__": + args = get_arg_parser().parse_args(sys.argv[1:]) + ui_thread(args.ip_address, args.frame_address) diff --git a/tools/replay/lib/__init__.py b/tools/replay/lib/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/tools/replay/lib/ui_helpers.py b/tools/replay/lib/ui_helpers.py new file mode 100644 index 0000000000..242e5a0c25 --- /dev/null +++ b/tools/replay/lib/ui_helpers.py @@ -0,0 +1,314 @@ +import platform +from collections import namedtuple + +import matplotlib +import matplotlib.pyplot as plt +import numpy as np +import pygame + +from tools.lib.lazy_property import lazy_property +from selfdrive.config import UIParams as UP +from selfdrive.config import RADAR_TO_CAMERA +from selfdrive.controls.lib.lane_planner import (compute_path_pinv, + model_polyfit) + +RED = (255, 0, 0) +GREEN = (0, 255, 0) +BLUE = (0, 0, 255) +YELLOW = (255, 255, 0) +BLACK = (0, 0, 0) +WHITE = (255, 255, 255) + +_PATH_X = np.arange(192.) +_PATH_XD = np.arange(192.) +_PATH_PINV = compute_path_pinv(50) +#_BB_OFFSET = 290, 332 +_BB_OFFSET = 0,0 +_BB_SCALE = 1164/640. +_BB_TO_FULL_FRAME = np.asarray([ + [_BB_SCALE, 0., _BB_OFFSET[0]], + [0., _BB_SCALE, _BB_OFFSET[1]], + [0., 0., 1.]]) +_FULL_FRAME_TO_BB = np.linalg.inv(_BB_TO_FULL_FRAME) + +METER_WIDTH = 20 + +ModelUIData = namedtuple("ModelUIData", ["cpath", "lpath", "rpath", "lead", "lead_future"]) + +_COLOR_CACHE = {} +def find_color(lidar_surface, color): + if color in _COLOR_CACHE: + return _COLOR_CACHE[color] + tcolor = 0 + ret = 255 + for x in lidar_surface.get_palette(): + #print tcolor, x + if x[0:3] == color: + ret = tcolor + break + tcolor += 1 + _COLOR_CACHE[color] = ret + return ret + +def warp_points(pt_s, warp_matrix): + # pt_s are the source points, nxm array. + pt_d = np.dot(warp_matrix[:, :-1], pt_s.T) + warp_matrix[:, -1, None] + + # Divide by last dimension for representation in image space. + return (pt_d[:-1, :] / pt_d[-1, :]).T + +def to_lid_pt(y, x): + px, py = -x * UP.lidar_zoom + UP.lidar_car_x, -y * UP.lidar_zoom + UP.lidar_car_y + if px > 0 and py > 0 and px < UP.lidar_x and py < UP.lidar_y: + return int(px), int(py) + return -1, -1 + + +def draw_path(y, x, color, img, calibration, top_down, lid_color=None): + # TODO: Remove big box. + uv_model_real = warp_points(np.column_stack((x, y)), calibration.car_to_model) + uv_model = np.round(uv_model_real).astype(int) + + uv_model_dots = uv_model[np.logical_and.reduce((np.all( # pylint: disable=no-member + uv_model > 0, axis=1), uv_model[:, 0] < img.shape[1] - 1, uv_model[:, 1] < + img.shape[0] - 1))] + + for i, j in ((-1, 0), (0, -1), (0, 0), (0, 1), (1, 0)): + img[uv_model_dots[:, 1] + i, uv_model_dots[:, 0] + j] = color + + # draw lidar path point on lidar + # find color in 8 bit + if lid_color is not None and top_down is not None: + tcolor = find_color(top_down[0], lid_color) + for i in range(len(x)): + px, py = to_lid_pt(x[i], y[i]) + if px != -1: + top_down[1][px, py] = tcolor + +def draw_steer_path(speed_ms, curvature, color, img, + calibration, top_down, VM, lid_color=None): + path_x = np.arange(101.) + path_y = np.multiply(path_x, np.tan(np.arcsin(np.clip(path_x * curvature, -0.999, 0.999)) / 2.)) + + draw_path(path_y, path_x, color, img, calibration, top_down, lid_color) + +def draw_lead_car(closest, top_down): + if closest != None: + closest_y = int(round(UP.lidar_car_y - closest * UP.lidar_zoom)) + if closest_y > 0: + top_down[1][int(round(UP.lidar_car_x - METER_WIDTH * 2)):int( + round(UP.lidar_car_x + METER_WIDTH * 2)), closest_y] = find_color( + top_down[0], (255, 0, 0)) + +def draw_lead_on(img, closest_x_m, closest_y_m, calibration, color, sz=10, img_offset=(0, 0)): + uv = warp_points(np.asarray([closest_x_m, closest_y_m]), calibration.car_to_bb)[0] + u, v = int(uv[0] + img_offset[0]), int(uv[1] + img_offset[1]) + if u > 0 and u < 640 and v > 0 and v < 480 - 5: + img[v - 5 - sz:v - 5 + sz, u] = color + img[v - 5, u - sz:u + sz] = color + return u, v + + +if platform.system() != 'Darwin': + matplotlib.use('QT4Agg') + + +def init_plots(arr, name_to_arr_idx, plot_xlims, plot_ylims, plot_names, plot_colors, plot_styles, bigplots=False): + color_palette = { "r": (1,0,0), + "g": (0,1,0), + "b": (0,0,1), + "k": (0,0,0), + "y": (1,1,0), + "p": (0,1,1), + "m": (1,0,1) } + + if bigplots == True: + fig = plt.figure(figsize=(6.4, 7.0)) + elif bigplots == False: + fig = plt.figure() + else: + fig = plt.figure(figsize=bigplots) + + fig.set_facecolor((0.2,0.2,0.2)) + + axs = [] + for pn in range(len(plot_ylims)): + ax = fig.add_subplot(len(plot_ylims),1,len(axs)+1) + ax.set_xlim(plot_xlims[pn][0], plot_xlims[pn][1]) + ax.set_ylim(plot_ylims[pn][0], plot_ylims[pn][1]) + ax.patch.set_facecolor((0.4, 0.4, 0.4)) + axs.append(ax) + + plots = [] ;idxs = [] ;plot_select = [] + for i, pl_list in enumerate(plot_names): + for j, item in enumerate(pl_list): + plot, = axs[i].plot(arr[:, name_to_arr_idx[item]], + label=item, + color=color_palette[plot_colors[i][j]], + linestyle=plot_styles[i][j]) + plots.append(plot) + idxs.append(name_to_arr_idx[item]) + plot_select.append(i) + axs[i].set_title(", ".join("%s (%s)" % (nm, cl) + for (nm, cl) in zip(pl_list, plot_colors[i])), fontsize=10) + if i < len(plot_ylims) - 1: + axs[i].set_xticks([]) + + fig.canvas.draw() + + renderer = fig.canvas.get_renderer() + + if matplotlib.get_backend() == "MacOSX": + fig.draw(renderer) + + def draw_plots(arr): + for ax in axs: + ax.draw_artist(ax.patch) + for i in range(len(plots)): + plots[i].set_ydata(arr[:, idxs[i]]) + axs[plot_select[i]].draw_artist(plots[i]) + + if matplotlib.get_backend() == "QT4Agg": + fig.canvas.update() + fig.canvas.flush_events() + + raw_data = renderer.tostring_rgb() + #print fig.canvas.get_width_height() + plot_surface = pygame.image.frombuffer(raw_data, fig.canvas.get_width_height(), "RGB").convert() + return plot_surface + + return draw_plots + + +def draw_mpc(liveMpc, top_down): + mpc_color = find_color(top_down[0], (0, 255, 0)) + for p in zip(liveMpc.x, liveMpc.y): + px, py = to_lid_pt(*p) + top_down[1][px, py] = mpc_color + + + +class CalibrationTransformsForWarpMatrix(object): + def __init__(self, model_to_full_frame, K, E): + self._model_to_full_frame = model_to_full_frame + self._K = K + self._E = E + + @property + def model_to_bb(self): + return _FULL_FRAME_TO_BB.dot(self._model_to_full_frame) + + @lazy_property + def model_to_full_frame(self): + return self._model_to_full_frame + + @lazy_property + def car_to_model(self): + return np.linalg.inv(self._model_to_full_frame).dot(self._K).dot( + self._E[:, [0, 1, 3]]) + + @lazy_property + def car_to_bb(self): + return _BB_TO_FULL_FRAME.dot(self._K).dot(self._E[:, [0, 1, 3]]) + + +def pygame_modules_have_loaded(): + return pygame.display.get_init() and pygame.font.get_init() + +def draw_var(y, x, var, color, img, calibration, top_down): + # otherwise drawing gets stupid + var = max(1e-1, min(var, 0.7)) + + varcolor = tuple(np.array(color)*0.5) + draw_path(y - var, x, varcolor, img, calibration, top_down) + draw_path(y + var, x, varcolor, img, calibration, top_down) + + +class ModelPoly(object): + def __init__(self, model_path): + if len(model_path.points) == 0 and len(model_path.poly) == 0: + self.valid = False + return + + if len(model_path.poly): + self.poly = np.array(model_path.poly) + else: + self.poly = model_polyfit(model_path.points, _PATH_PINV) + + self.prob = model_path.prob + self.std = model_path.std + self.y = np.polyval(self.poly, _PATH_XD) + self.valid = True + +def extract_model_data(md): + return ModelUIData( + cpath=ModelPoly(md.path), + lpath=ModelPoly(md.leftLane), + rpath=ModelPoly(md.rightLane), + lead=md.lead, + lead_future=md.leadFuture, + ) + +def plot_model(m, VM, v_ego, curvature, imgw, calibration, top_down, d_poly, top_down_color=216): + if calibration is None or top_down is None: + return + + for lead in [m.lead, m.lead_future]: + if lead.prob < 0.5: + continue + + lead_dist_from_radar = lead.dist - RADAR_TO_CAMERA + _, py_top = to_lid_pt(lead_dist_from_radar + lead.std, lead.relY) + px, py_bottom = to_lid_pt(lead_dist_from_radar - lead.std, lead.relY) + top_down[1][int(round(px - 4)):int(round(px + 4)), py_top:py_bottom] = top_down_color + + color = (0, int(255 * m.lpath.prob), 0) + for path in [m.cpath, m.lpath, m.rpath]: + if path.valid: + draw_path(path.y, _PATH_XD, color, imgw, calibration, top_down, YELLOW) + draw_var(path.y, _PATH_XD, path.std, color, imgw, calibration, top_down) + + if d_poly is not None: + dpath_y = np.polyval(d_poly, _PATH_X) + draw_path(dpath_y, _PATH_X, RED, imgw, calibration, top_down, RED) + + # draw user path from curvature + draw_steer_path(v_ego, curvature, BLUE, imgw, calibration, top_down, VM, BLUE) + + +def maybe_update_radar_points(lt, lid_overlay): + ar_pts = [] + if lt is not None: + ar_pts = {} + for track in lt: + ar_pts[track.trackId] = [track.dRel, track.yRel, track.vRel, track.aRel, track.oncoming, track.stationary] + for ids, pt in ar_pts.items(): + px, py = to_lid_pt(pt[0], pt[1]) + if px != -1: + if pt[-1]: + color = 240 + elif pt[-2]: + color = 230 + else: + color = 255 + if int(ids) == 1: + lid_overlay[px - 2:px + 2, py - 10:py + 10] = 100 + else: + lid_overlay[px - 2:px + 2, py - 2:py + 2] = color + +def get_blank_lid_overlay(UP): + lid_overlay = np.zeros((UP.lidar_x, UP.lidar_y), 'uint8') + # Draw the car. + lid_overlay[int(round(UP.lidar_car_x - UP.car_hwidth)):int( + round(UP.lidar_car_x + UP.car_hwidth)), int(round(UP.lidar_car_y - + UP.car_front))] = UP.car_color + lid_overlay[int(round(UP.lidar_car_x - UP.car_hwidth)):int( + round(UP.lidar_car_x + UP.car_hwidth)), int(round(UP.lidar_car_y + + UP.car_back))] = UP.car_color + lid_overlay[int(round(UP.lidar_car_x - UP.car_hwidth)), int( + round(UP.lidar_car_y - UP.car_front)):int(round( + UP.lidar_car_y + UP.car_back))] = UP.car_color + lid_overlay[int(round(UP.lidar_car_x + UP.car_hwidth)), int( + round(UP.lidar_car_y - UP.car_front)):int(round( + UP.lidar_car_y + UP.car_back))] = UP.car_color + return lid_overlay diff --git a/tools/replay/mapd.py b/tools/replay/mapd.py new file mode 100755 index 0000000000..9b5e0bccd5 --- /dev/null +++ b/tools/replay/mapd.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python + +import matplotlib +matplotlib.use('TkAgg') +import matplotlib.pyplot as plt + +import numpy as np +import zmq +from cereal.services import service_list +from selfdrive.config import Conversions as CV +import cereal.messaging as messaging + + +if __name__ == "__main__": + live_map_sock = messaging.sub_sock(service_list['liveMapData'].port, conflate=True) + plan_sock = messaging.sub_sock(service_list['plan'].port, conflate=True) + + plt.ion() + fig = plt.figure(figsize=(8, 16)) + ax = fig.add_subplot(2, 1, 1) + ax.set_title('Map') + + SCALE = 1000 + ax.set_xlim([-SCALE, SCALE]) + ax.set_ylim([-SCALE, SCALE]) + ax.set_xlabel('x [m]') + ax.set_ylabel('y [m]') + ax.grid(True) + + points_plt, = ax.plot([0.0], [0.0], "--xk") + cur, = ax.plot([0.0], [0.0], "xr") + + speed_txt = ax.text(-500, 900, '') + curv_txt = ax.text(-500, 775, '') + + ax = fig.add_subplot(2, 1, 2) + ax.set_title('Curvature') + curvature_plt, = ax.plot([0.0], [0.0], "--xk") + ax.set_xlim([0, 500]) + ax.set_ylim([0, 1e-2]) + ax.set_xlabel('Distance along path [m]') + ax.set_ylabel('Curvature [1/m]') + ax.grid(True) + + plt.show() + + while True: + m = messaging.recv_one_or_none(live_map_sock) + p = messaging.recv_one_or_none(plan_sock) + if p is not None: + v = p.plan.vCurvature * CV.MS_TO_MPH + speed_txt.set_text('Desired curvature speed: %.2f mph' % v) + + if m is not None: + print("Current way id: %d" % m.liveMapData.wayId) + curv_txt.set_text('Curvature valid: %s Dist: %03.0f m\nSpeedlimit valid: %s Speed: %.0f mph' % + (str(m.liveMapData.curvatureValid), + m.liveMapData.distToTurn, + str(m.liveMapData.speedLimitValid), + m.liveMapData.speedLimit * CV.MS_TO_MPH)) + + points_plt.set_xdata(m.liveMapData.roadX) + points_plt.set_ydata(m.liveMapData.roadY) + curvature_plt.set_xdata(m.liveMapData.roadCurvatureX) + curvature_plt.set_ydata(m.liveMapData.roadCurvature) + + fig.canvas.draw() + fig.canvas.flush_events() diff --git a/tools/replay/rqplot.py b/tools/replay/rqplot.py new file mode 100755 index 0000000000..32e9a34981 --- /dev/null +++ b/tools/replay/rqplot.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python +import os +import sys +import matplotlib.pyplot as plt +import numpy as np +import cereal.messaging as messaging +import time + + +# tool to plot one or more signals live. Call ex: +#./rqplot.py log.carState.vEgo log.carState.aEgo + +# TODO: can this tool consume 10x less cpu? + +def recursive_getattr(x, name): + l = name.split('.') + if len(l) == 1: + return getattr(x, name) + else: + return recursive_getattr(getattr(x, l[0]), ".".join(l[1:]) ) + + +if __name__ == "__main__": + poller = messaging.Poller() + + services = [] + fields = [] + subs = [] + values = [] + + plt.ion() + fig, ax = plt.subplots() + #fig = plt.figure(figsize=(10, 15)) + #ax = fig.add_subplot(111) + ax.grid(True) + fig.canvas.draw() + + subs_name = sys.argv[1:] + lines = [] + x, y = [], [] + LEN = 500 + + for i, sub in enumerate(subs_name): + sub_split = sub.split(".") + services.append(sub_split[0]) + fields.append(".".join(sub_split[1:])) + subs.append(messaging.sub_sock(sub_split[0], poller)) + + x.append(np.ones(LEN)*np.nan) + y.append(np.ones(LEN)*np.nan) + lines.append(ax.plot(x[i], y[i])[0]) + + for l in lines: + l.set_marker("*") + + cur_t = 0. + ax.legend(subs_name) + ax.set_xlabel('time [s]') + + while 1: + print(1./(time.time() - cur_t)) + cur_t = time.time() + for i, s in enumerate(subs): + msg = messaging.recv_sock(s) + #msg = messaging.recv_one_or_none(s) + if msg is not None: + x[i] = np.append(x[i], getattr(msg, 'logMonoTime') / float(1e9)) + x[i] = np.delete(x[i], 0) + y[i] = np.append(y[i], recursive_getattr(msg, subs_name[i])) + y[i] = np.delete(y[i], 0) + + lines[i].set_xdata(x[i]) + lines[i].set_ydata(y[i]) + + ax.relim() + ax.autoscale_view(True, scaley=True, scalex=True) + + fig.canvas.blit(ax.bbox) + fig.canvas.flush_events() + + # just a bit of wait to avoid 100% CPU usage + time.sleep(0.001) + diff --git a/tools/replay/ui.py b/tools/replay/ui.py new file mode 100755 index 0000000000..0dfa36ac98 --- /dev/null +++ b/tools/replay/ui.py @@ -0,0 +1,278 @@ +#!/usr/bin/env python +import argparse +import os +import sys + +os.environ["OMP_NUM_THREADS"] = "1" + +import cv2 +import numpy as np +import pygame + +from common.basedir import BASEDIR +from common.transformations.camera import FULL_FRAME_SIZE, eon_intrinsics +from common.transformations.model import (MODEL_CX, MODEL_CY, MODEL_INPUT_SIZE, + get_camera_frame_from_model_frame) +from selfdrive.car.toyota.interface import CarInterface as ToyotaInterface +from selfdrive.config import UIParams as UP +from selfdrive.controls.lib.vehicle_model import VehicleModel +import cereal.messaging as messaging +from tools.replay.lib.ui_helpers import (_BB_TO_FULL_FRAME, BLACK, BLUE, GREEN, + YELLOW, RED, + CalibrationTransformsForWarpMatrix, + draw_lead_car, draw_lead_on, draw_mpc, + extract_model_data, + get_blank_lid_overlay, init_plots, + maybe_update_radar_points, plot_model, + pygame_modules_have_loaded, + warp_points) + +os.environ['BASEDIR'] = BASEDIR + +ANGLE_SCALE = 5.0 +HOR = os.getenv("HORIZONTAL") is not None + + +def ui_thread(addr, frame_address): + # TODO: Detect car from replay and use that to select carparams + CP = ToyotaInterface.get_params("TOYOTA PRIUS 2017") + VM = VehicleModel(CP) + + CalP = np.asarray([[0, 0], [MODEL_INPUT_SIZE[0], 0], [MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]], [0, MODEL_INPUT_SIZE[1]]]) + vanishing_point = np.asarray([[MODEL_CX, MODEL_CY]]) + + pygame.init() + pygame.font.init() + assert pygame_modules_have_loaded() + + if HOR: + size = (640+384+640, 960) + write_x = 5 + write_y = 680 + else: + size = (640+384, 960+300) + write_x = 645 + write_y = 970 + + pygame.display.set_caption("openpilot debug UI") + screen = pygame.display.set_mode(size, pygame.DOUBLEBUF) + + alert1_font = pygame.font.SysFont("arial", 30) + alert2_font = pygame.font.SysFont("arial", 20) + info_font = pygame.font.SysFont("arial", 15) + + camera_surface = pygame.surface.Surface((640, 480), 0, 24).convert() + cameraw_surface = pygame.surface.Surface(MODEL_INPUT_SIZE, 0, 24).convert() + cameraw_test_surface = pygame.surface.Surface(MODEL_INPUT_SIZE, 0, 24) + top_down_surface = pygame.surface.Surface((UP.lidar_x, UP.lidar_y),0,8) + + frame = messaging.sub_sock('frame', addr=addr, conflate=True) + sm = messaging.SubMaster(['carState', 'plan', 'carControl', 'radarState', 'liveCalibration', 'controlsState', 'liveTracks', 'model', 'liveMpc', 'liveParameters', 'pathPlan'], addr=addr) + + calibration = None + img = np.zeros((480, 640, 3), dtype='uint8') + imgff = np.zeros((FULL_FRAME_SIZE[1], FULL_FRAME_SIZE[0], 3), dtype=np.uint8) + imgw = np.zeros((160, 320, 3), dtype=np.uint8) # warped image + lid_overlay_blank = get_blank_lid_overlay(UP) + + # plots + name_to_arr_idx = { "gas": 0, + "computer_gas": 1, + "user_brake": 2, + "computer_brake": 3, + "v_ego": 4, + "v_pid": 5, + "angle_steers_des": 6, + "angle_steers": 7, + "angle_steers_k": 8, + "steer_torque": 9, + "v_override": 10, + "v_cruise": 11, + "a_ego": 12, + "a_target": 13, + "accel_override": 14} + + plot_arr = np.zeros((100, len(name_to_arr_idx.values()))) + + plot_xlims = [(0, plot_arr.shape[0]), (0, plot_arr.shape[0]), (0, plot_arr.shape[0]), (0, plot_arr.shape[0])] + plot_ylims = [(-0.1, 1.1), (-ANGLE_SCALE, ANGLE_SCALE), (0., 75.), (-3.0, 2.0)] + plot_names = [["gas", "computer_gas", "user_brake", "computer_brake", "accel_override"], + ["angle_steers", "angle_steers_des", "angle_steers_k", "steer_torque"], + ["v_ego", "v_override", "v_pid", "v_cruise"], + ["a_ego", "a_target"]] + plot_colors = [["b", "b", "g", "r", "y"], + ["b", "g", "y", "r"], + ["b", "g", "r", "y"], + ["b", "r"]] + plot_styles = [["-", "-", "-", "-", "-"], + ["-", "-", "-", "-"], + ["-", "-", "-", "-"], + ["-", "-"]] + + draw_plots = init_plots(plot_arr, name_to_arr_idx, plot_xlims, plot_ylims, plot_names, plot_colors, plot_styles, bigplots=True) + + counter = 0 + while 1: + list(pygame.event.get()) + + screen.fill((64,64,64)) + lid_overlay = lid_overlay_blank.copy() + top_down = top_down_surface, lid_overlay + + # ***** frame ***** + fpkt = messaging.recv_one(frame) + rgb_img_raw = fpkt.frame.image + + if fpkt.frame.transform: + img_transform = np.array(fpkt.frame.transform).reshape(3,3) + else: + # assume frame is flipped + img_transform = np.array([ + [-1.0, 0.0, FULL_FRAME_SIZE[0]-1], + [ 0.0, -1.0, FULL_FRAME_SIZE[1]-1], + [ 0.0, 0.0, 1.0] + ]) + + + if rgb_img_raw and len(rgb_img_raw) == FULL_FRAME_SIZE[0] * FULL_FRAME_SIZE[1] * 3: + imgff = np.frombuffer(rgb_img_raw, dtype=np.uint8).reshape((FULL_FRAME_SIZE[1], FULL_FRAME_SIZE[0], 3)) + imgff = imgff[:, :, ::-1] # Convert BGR to RGB + cv2.warpAffine(imgff, np.dot(img_transform, _BB_TO_FULL_FRAME)[:2], + (img.shape[1], img.shape[0]), dst=img, flags=cv2.WARP_INVERSE_MAP) + + intrinsic_matrix = eon_intrinsics + else: + img.fill(0) + intrinsic_matrix = np.eye(3) + + if calibration is not None: + transform = np.dot(img_transform, calibration.model_to_full_frame) + imgw = cv2.warpAffine(imgff, transform[:2], (MODEL_INPUT_SIZE[0], MODEL_INPUT_SIZE[1]), flags=cv2.WARP_INVERSE_MAP) + else: + imgw.fill(0) + + sm.update() + + w = sm['controlsState'].lateralControlState.which() + if w == 'lqrState': + angle_steers_k = sm['controlsState'].lateralControlState.lqrState.steerAngle + elif w == 'indiState': + angle_steers_k = sm['controlsState'].lateralControlState.indiState.steerAngle + else: + angle_steers_k = np.inf + + plot_arr[:-1] = plot_arr[1:] + plot_arr[-1, name_to_arr_idx['angle_steers']] = sm['controlsState'].angleSteers + plot_arr[-1, name_to_arr_idx['angle_steers_des']] = sm['carControl'].actuators.steerAngle + plot_arr[-1, name_to_arr_idx['angle_steers_k']] = angle_steers_k + plot_arr[-1, name_to_arr_idx['gas']] = sm['carState'].gas + plot_arr[-1, name_to_arr_idx['computer_gas']] = sm['carControl'].actuators.gas + plot_arr[-1, name_to_arr_idx['user_brake']] = sm['carState'].brake + plot_arr[-1, name_to_arr_idx['steer_torque']] = sm['carControl'].actuators.steer * ANGLE_SCALE + plot_arr[-1, name_to_arr_idx['computer_brake']] = sm['carControl'].actuators.brake + plot_arr[-1, name_to_arr_idx['v_ego']] = sm['controlsState'].vEgo + plot_arr[-1, name_to_arr_idx['v_pid']] = sm['controlsState'].vPid + plot_arr[-1, name_to_arr_idx['v_override']] = sm['carControl'].cruiseControl.speedOverride + plot_arr[-1, name_to_arr_idx['v_cruise']] = sm['carState'].cruiseState.speed + plot_arr[-1, name_to_arr_idx['a_ego']] = sm['carState'].aEgo + plot_arr[-1, name_to_arr_idx['a_target']] = sm['plan'].aTarget + plot_arr[-1, name_to_arr_idx['accel_override']] = sm['carControl'].cruiseControl.accelOverride + + # ***** model **** + if len(sm['model'].path.poly) > 0: + model_data = extract_model_data(sm['model']) + plot_model(model_data, VM, sm['controlsState'].vEgo, sm['controlsState'].curvature, imgw, calibration, + top_down, np.array(sm['pathPlan'].dPoly)) + + # MPC + if sm.updated['liveMpc']: + draw_mpc(sm['liveMpc'], top_down) + + # draw all radar points + maybe_update_radar_points(sm['liveTracks'], top_down[1]) + + + if sm.updated['liveCalibration']: + extrinsic_matrix = np.asarray(sm['liveCalibration'].extrinsicMatrix).reshape(3, 4) + ke = intrinsic_matrix.dot(extrinsic_matrix) + warp_matrix = get_camera_frame_from_model_frame(ke) + calibration = CalibrationTransformsForWarpMatrix(warp_matrix, intrinsic_matrix, extrinsic_matrix) + + # draw red pt for lead car in the main img + for lead in [sm['radarState'].leadOne, sm['radarState'].leadTwo]: + if lead.status: + if calibration is not None: + draw_lead_on(img, lead.dRel, lead.yRel, calibration, color=(192,0,0)) + + draw_lead_car(lead.dRel, top_down) + + # *** blits *** + pygame.surfarray.blit_array(camera_surface, img.swapaxes(0,1)) + screen.blit(camera_surface, (0, 0)) + + # display alerts + alert_line1 = alert1_font.render(sm['controlsState'].alertText1, True, (255,0,0)) + alert_line2 = alert2_font.render(sm['controlsState'].alertText2, True, (255,0,0)) + screen.blit(alert_line1, (180, 150)) + screen.blit(alert_line2, (180, 190)) + + if calibration is not None and img is not None: + cpw = warp_points(CalP, calibration.model_to_bb) + vanishing_pointw = warp_points(vanishing_point, calibration.model_to_bb) + pygame.draw.polygon(screen, BLUE, tuple(map(tuple, cpw)), 1) + pygame.draw.circle(screen, BLUE, list(map(int, map(round, vanishing_pointw[0]))), 2) + + if HOR: + screen.blit(draw_plots(plot_arr), (640+384, 0)) + else: + screen.blit(draw_plots(plot_arr), (0, 600)) + + pygame.surfarray.blit_array(cameraw_surface, imgw.swapaxes(0, 1)) + screen.blit(cameraw_surface, (320, 480)) + + pygame.surfarray.blit_array(*top_down) + screen.blit(top_down[0], (640,0)) + + i = 0 + SPACING = 25 + + lines = [ + info_font.render("ENABLED", True, GREEN if sm['controlsState'].enabled else BLACK), + info_font.render("BRAKE LIGHTS", True, RED if sm['carState'].brakeLights else BLACK), + info_font.render("SPEED: " + str(round(sm['carState'].vEgo, 1)) + " m/s", True, YELLOW), + info_font.render("LONG CONTROL STATE: " + str(sm['controlsState'].longControlState), True, YELLOW), + info_font.render("LONG MPC SOURCE: " + str(sm['plan'].longitudinalPlanSource), True, YELLOW), + None, + info_font.render("ANGLE OFFSET (AVG): " + str(round(sm['liveParameters'].angleOffsetAverage, 2)) + " deg", True, YELLOW), + info_font.render("ANGLE OFFSET (INSTANT): " + str(round(sm['liveParameters'].angleOffset, 2)) + " deg", True, YELLOW), + info_font.render("STIFFNESS: " + str(round(sm['liveParameters'].stiffnessFactor * 100., 2)) + " %", True, YELLOW), + info_font.render("STEER RATIO: " + str(round(sm['liveParameters'].steerRatio, 2)), True, YELLOW) + ] + + for i, line in enumerate(lines): + if line is not None: + screen.blit(line, (write_x, write_y + i * SPACING)) + + # this takes time...vsync or something + pygame.display.flip() + +def get_arg_parser(): + parser = argparse.ArgumentParser( + description="Show replay data in a UI.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + parser.add_argument("ip_address", nargs="?", default="127.0.0.1", + help="The ip address on which to receive zmq messages.") + + parser.add_argument("--frame-address", default=None, + help="The frame address (fully qualified ZMQ endpoint for frames) on which to receive zmq messages.") + return parser + +if __name__ == "__main__": + args = get_arg_parser().parse_args(sys.argv[1:]) + + if args.ip_address != "127.0.0.1": + os.environ["ZMQ"] = "1" + messaging.context = messaging.Context() + + ui_thread(args.ip_address, args.frame_address) diff --git a/tools/replay/unlogger.py b/tools/replay/unlogger.py new file mode 100755 index 0000000000..8208df7a44 --- /dev/null +++ b/tools/replay/unlogger.py @@ -0,0 +1,446 @@ +#!/usr/bin/env python +import argparse +import os +import sys +import zmq +import time +import gc +import signal +from threading import Thread +import numpy as np +from uuid import uuid4 +from collections import namedtuple +from collections import deque +from multiprocessing import Process, TimeoutError +from datetime import datetime + +# strat 1: script to copy files +# strat 2: build pip packages around these +# could be its own pip package, which we'd need to build and release +from cereal import log as capnp_log +from cereal.services import service_list +from cereal.messaging import pub_sock, MultiplePublishersError +from common import realtime + +from tools.lib.file_helpers import mkdirs_exists_ok +from tools.lib.kbhit import KBHit +from tools.lib.logreader import MultiLogIterator +from tools.lib.route import Route +from tools.lib.route_framereader import RouteFrameReader + +# Commands. +SetRoute = namedtuple("SetRoute", ("name", "start_time", "data_dir")) +SeekAbsoluteTime = namedtuple("SeekAbsoluteTime", ("secs",)) +SeekRelativeTime = namedtuple("SeekRelativeTime", ("secs",)) +TogglePause = namedtuple("TogglePause", ()) +StopAndQuit = namedtuple("StopAndQuit", ()) + + +class UnloggerWorker(object): + def __init__(self): + self._frame_reader = None + self._cookie = None + self._readahead = deque() + + def run(self, commands_address, data_address, pub_types): + zmq.Context._instance = None + commands_socket = zmq.Context.instance().socket(zmq.PULL) + commands_socket.connect(commands_address) + + data_socket = zmq.Context.instance().socket(zmq.PUSH) + data_socket.connect(data_address) + + poller = zmq.Poller() + poller.register(commands_socket, zmq.POLLIN) + + # We can't publish frames without encodeIdx, so add when it's missing. + if "frame" in pub_types: + pub_types["encodeIdx"] = None + + # gc.set_debug(gc.DEBUG_LEAK | gc.DEBUG_OBJECTS | gc.DEBUG_STATS | gc.DEBUG_SAVEALL | + # gc.DEBUG_UNCOLLECTABLE) + + # TODO: WARNING pycapnp leaks memory all over the place after unlogger runs for a while, gc + # pauses become huge because there are so many tracked objects solution will be to switch to new + # cython capnp + try: + route = None + while True: + while poller.poll(0.) or route is None: + cookie, cmd = commands_socket.recv_pyobj() + route = self._process_commands(cmd, route) + + # **** get message **** + self._read_logs(cookie, pub_types) + self._send_logs(data_socket) + finally: + if self._frame_reader is not None: + self._frame_reader.close() + data_socket.close() + commands_socket.close() + + def _read_logs(self, cookie, pub_types): + fullHEVC = capnp_log.EncodeIndex.Type.fullHEVC + lr = self._lr + while len(self._readahead) < 1000: + route_time = lr.tell() + msg = next(lr) + typ = msg.which() + if typ not in pub_types: + continue + + # **** special case certain message types **** + if typ == "encodeIdx" and msg.encodeIdx.type == fullHEVC: + # this assumes the encodeIdx always comes before the frame + self._frame_id_lookup[ + msg.encodeIdx.frameId] = msg.encodeIdx.segmentNum, msg.encodeIdx.segmentId + #print "encode", msg.encodeIdx.frameId, len(self._readahead), route_time + self._readahead.appendleft((typ, msg, route_time, cookie)) + + def _send_logs(self, data_socket): + while len(self._readahead) > 500: + typ, msg, route_time, cookie = self._readahead.pop() + smsg = msg.as_builder() + + if typ == "frame": + frame_id = msg.frame.frameId + + # Frame exists, make sure we have a framereader. + # load the frame readers as needed + s1 = time.time() + img = self._frame_reader.get(frame_id, pix_fmt="rgb24") + fr_time = time.time() - s1 + if fr_time > 0.05: + print("FRAME(%d) LAG -- %.2f ms" % (frame_id, fr_time*1000.0)) + + if img is not None: + img = img[:, :, ::-1] # Convert RGB to BGR, which is what the camera outputs + img = img.flatten() + smsg.frame.image = img.tobytes() + + data_socket.send_pyobj((cookie, typ, msg.logMonoTime, route_time), flags=zmq.SNDMORE) + data_socket.send(smsg.to_bytes(), copy=False) + + def _process_commands(self, cmd, route): + seek_to = None + if route is None or (isinstance(cmd, SetRoute) and route.name != cmd.name): + seek_to = cmd.start_time + route = Route(cmd.name, cmd.data_dir) + self._lr = MultiLogIterator(route.log_paths(), wraparound=True) + if self._frame_reader is not None: + self._frame_reader.close() + # reset frames for a route + self._frame_id_lookup = {} + self._frame_reader = RouteFrameReader( + route.camera_paths(), None, self._frame_id_lookup, readahead=True) + + # always reset this on a seek + if isinstance(cmd, SeekRelativeTime): + seek_to = self._lr.tell() + cmd.secs + elif isinstance(cmd, SeekAbsoluteTime): + seek_to = cmd.secs + elif isinstance(cmd, StopAndQuit): + exit() + + if seek_to is not None: + print("seeking", seek_to) + if not self._lr.seek(seek_to): + print("Can't seek: time out of bounds") + else: + next(self._lr) # ignore one + return route + +def _get_address_send_func(address): + sock = pub_sock(address) + return sock.send + + +def unlogger_thread(command_address, forward_commands_address, data_address, run_realtime, + address_mapping, publish_time_length, bind_early, no_loop): + # Clear context to avoid problems with multiprocessing. + zmq.Context._instance = None + context = zmq.Context.instance() + + command_sock = context.socket(zmq.PULL) + command_sock.bind(command_address) + + forward_commands_socket = context.socket(zmq.PUSH) + forward_commands_socket.bind(forward_commands_address) + + data_socket = context.socket(zmq.PULL) + data_socket.bind(data_address) + + # Set readahead to a reasonable number. + data_socket.setsockopt(zmq.RCVHWM, 10000) + + poller = zmq.Poller() + poller.register(command_sock, zmq.POLLIN) + poller.register(data_socket, zmq.POLLIN) + + if bind_early: + send_funcs = { + typ: _get_address_send_func(address) + for typ, address in address_mapping.items() + } + + # Give subscribers a chance to connect. + time.sleep(0.1) + else: + send_funcs = {} + + start_time = float("inf") + printed_at = 0 + generation = 0 + paused = False + reset_time = True + prev_msg_time = None + while True: + evts = dict(poller.poll()) + if command_sock in evts: + cmd = command_sock.recv_pyobj() + if isinstance(cmd, TogglePause): + paused = not paused + if paused: + poller.modify(data_socket, 0) + else: + poller.modify(data_socket, zmq.POLLIN) + else: + # Forward the command the the log data thread. + # TODO: Remove everything on data_socket. + generation += 1 + forward_commands_socket.send_pyobj((generation, cmd)) + if isinstance(cmd, StopAndQuit): + return + + reset_time = True + elif data_socket in evts: + msg_generation, typ, msg_time, route_time = data_socket.recv_pyobj(flags=zmq.RCVMORE) + msg_bytes = data_socket.recv() + if msg_generation < generation: + # Skip packets. + continue + + if no_loop and prev_msg_time is not None and prev_msg_time > msg_time + 1e9: + generation += 1 + forward_commands_socket.send_pyobj((generation, StopAndQuit())) + return + prev_msg_time = msg_time + + msg_time_seconds = msg_time * 1e-9 + if reset_time: + msg_start_time = msg_time_seconds + real_start_time = realtime.sec_since_boot() + start_time = min(start_time, msg_start_time) + reset_time = False + + if publish_time_length and msg_time_seconds - start_time > publish_time_length: + generation += 1 + forward_commands_socket.send_pyobj((generation, StopAndQuit())) + return + + # Print time. + if abs(printed_at - route_time) > 5.: + print("at", route_time) + printed_at = route_time + + if typ not in send_funcs: + if typ in address_mapping: + # Remove so we don't keep printing warnings. + address = address_mapping.pop(typ) + try: + print("binding", typ) + send_funcs[typ] = _get_address_send_func(address) + except Exception as e: + print("couldn't replay {}: {}".format(typ, e)) + continue + else: + # Skip messages that we are not registered to publish. + continue + + # Sleep as needed for real time playback. + if run_realtime: + msg_time_offset = msg_time_seconds - msg_start_time + real_time_offset = realtime.sec_since_boot() - real_start_time + lag = msg_time_offset - real_time_offset + if lag > 0 and lag < 30: # a large jump is OK, likely due to an out of order segment + if lag > 1: + print("sleeping for", lag) + time.sleep(lag) + elif lag < -1: + # Relax the real time schedule when we slip far behind. + reset_time = True + + # Send message. + try: + send_funcs[typ](msg_bytes) + except MultiplePublishersError: + del send_funcs[typ] + +def timestamp_to_s(tss): + return time.mktime(datetime.strptime(tss, '%Y-%m-%d--%H-%M-%S').timetuple()) + +def absolute_time_str(s, start_time): + try: + # first try if it's a float + return float(s) + except ValueError: + # now see if it's a timestamp + return timestamp_to_s(s) - start_time + +def _get_address_mapping(args): + if args.min is not None: + services_to_mock = [ + 'thermal', 'can', 'health', 'sensorEvents', 'gpsNMEA', 'frame', 'encodeIdx', + 'model', 'features', 'liveLocation', 'gpsLocation' + ] + elif args.enabled is not None: + services_to_mock = args.enabled + else: + services_to_mock = service_list.keys() + + address_mapping = {service_name: service_name for service_name in services_to_mock} + address_mapping.update(dict(args.address_mapping)) + + for k in args.disabled: + address_mapping.pop(k, None) + + non_services = set(address_mapping) - set(service_list) + if non_services: + print("WARNING: Unknown services {}".format(list(non_services))) + + return address_mapping + +def keyboard_controller_thread(q, route_start_time): + print("keyboard waiting for input") + kb = KBHit() + while 1: + c = kb.getch() + if c=='m': # Move forward by 1m + q.send_pyobj(SeekRelativeTime(60)) + elif c=='M': # Move backward by 1m + q.send_pyobj(SeekRelativeTime(-60)) + elif c=='s': # Move forward by 10s + q.send_pyobj(SeekRelativeTime(10)) + elif c=='S': # Move backward by 10s + q.send_pyobj(SeekRelativeTime(-10)) + elif c=='G': # Move backward by 10s + q.send_pyobj(SeekAbsoluteTime(0.)) + elif c=="\x20": # Space bar. + q.send_pyobj(TogglePause()) + elif c=="\n": + try: + seek_time_input = raw_input('time: ') + seek_time = absolute_time_str(seek_time_input, route_start_time) + + q.send_pyobj(SeekAbsoluteTime(seek_time)) + except Exception as e: + print("Time not understood: {}".format(e)) + +def get_arg_parser(): + parser = argparse.ArgumentParser( + description="Mock openpilot components by publishing logged messages.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter) + + parser.add_argument("route_name", type=(lambda x: x.replace("#", "|")), nargs="?", + help="The route whose messages will be published.") + parser.add_argument("data_dir", nargs='?', default=os.getenv('UNLOGGER_DATA_DIR'), + help="Path to directory in which log and camera files are located.") + + parser.add_argument("--no-loop", action="store_true", help="Stop at the end of the replay.") + + key_value_pair = lambda x: x.split("=") + parser.add_argument("address_mapping", nargs="*", type=key_value_pair, + help="Pairs = to publish on .") + + comma_list = lambda x: x.split(",") + to_mock_group = parser.add_mutually_exclusive_group() + to_mock_group.add_argument("--min", action="store_true", default=os.getenv("MIN")) + to_mock_group.add_argument("--enabled", default=os.getenv("ENABLED"), type=comma_list) + + parser.add_argument("--disabled", type=comma_list, default=os.getenv("DISABLED") or ()) + + parser.add_argument( + "--tl", dest="publish_time_length", type=float, default=None, + help="Length of interval in event time for which messages should be published.") + + parser.add_argument( + "--no-realtime", dest="realtime", action="store_false", default=True, + help="Publish messages as quickly as possible instead of realtime.") + + parser.add_argument( + "--no-interactive", dest="interactive", action="store_false", default=True, + help="Disable interactivity.") + + parser.add_argument( + "--bind-early", action="store_true", default=False, + help="Bind early to avoid dropping messages.") + + return parser + +def main(argv): + args = get_arg_parser().parse_args(sys.argv[1:]) + + command_address = "ipc:///tmp/{}".format(uuid4()) + forward_commands_address = "ipc:///tmp/{}".format(uuid4()) + data_address = "ipc:///tmp/{}".format(uuid4()) + + address_mapping = _get_address_mapping(args) + + command_sock = zmq.Context.instance().socket(zmq.PUSH) + command_sock.connect(command_address) + + if args.route_name is not None: + route_name_split = args.route_name.split("|") + if len(route_name_split) > 1: + route_start_time = timestamp_to_s(route_name_split[1]) + else: + route_start_time = 0 + command_sock.send_pyobj( + SetRoute(args.route_name, 0, args.data_dir)) + else: + print("waiting for external command...") + route_start_time = 0 + + subprocesses = {} + try: + subprocesses["data"] = Process( + target=UnloggerWorker().run, + args=(forward_commands_address, data_address, address_mapping.copy())) + + subprocesses["control"] = Process( + target=unlogger_thread, + args=(command_address, forward_commands_address, data_address, args.realtime, + _get_address_mapping(args), args.publish_time_length, args.bind_early, args.no_loop)) + + for p in subprocesses.values(): + p.daemon = True + + subprocesses["data"].start() + subprocesses["control"].start() + + # Exit if any of the children die. + def exit_if_children_dead(*_): + for name, p in subprocesses.items(): + if not p.is_alive(): + [p.terminate() for p in subprocesses.values()] + exit() + signal.signal(signal.SIGCHLD, signal.SIGIGN) + signal.signal(signal.SIGCHLD, exit_if_children_dead) + + if args.interactive: + keyboard_controller_thread(command_sock, route_start_time) + else: + # Wait forever for children. + while True: + time.sleep(10000.) + finally: + for p in subprocesses.values(): + if p.is_alive(): + try: + p.join(3.) + except TimeoutError: + p.terminate() + continue + +if __name__ == "__main__": + sys.exit(main(sys.argv[1:])) diff --git a/tools/requirements.txt b/tools/requirements.txt new file mode 100644 index 0000000000..ef1e743910 --- /dev/null +++ b/tools/requirements.txt @@ -0,0 +1,11 @@ +aenum +atomicwrites +futures +libarchive +lru-dict +matplotlib==2.0.2 +numpy +opencv-python +pygame +hexdump==3.3 +av==0.5.0 diff --git a/tools/sim/.gitignore b/tools/sim/.gitignore new file mode 100644 index 0000000000..486a0508d3 --- /dev/null +++ b/tools/sim/.gitignore @@ -0,0 +1,3 @@ +CARLA_*.tar.gz +carla + diff --git a/tools/sim/can.py b/tools/sim/can.py new file mode 100755 index 0000000000..be806ff276 --- /dev/null +++ b/tools/sim/can.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +import time +import cereal.messaging as messaging +from opendbc.can.parser import CANParser +from opendbc.can.packer import CANPacker +from selfdrive.boardd.boardd_api_impl import can_list_to_can_capnp +from selfdrive.car.honda.values import FINGERPRINTS, CAR +from selfdrive.car import crc8_pedal + +from selfdrive.test.longitudinal_maneuvers.plant import get_car_can_parser +cp = get_car_can_parser() +#cp = CANParser("honda_civic_touring_2016_can_generated") + +packer = CANPacker("honda_civic_touring_2016_can_generated") +rpacker = CANPacker("acura_ilx_2016_nidec") + +def can_function(pm, speed, angle, idx, engage): + msg = [] + msg.append(packer.make_can_msg("ENGINE_DATA", 0, {"XMISSION_SPEED": speed}, idx)) + msg.append(packer.make_can_msg("WHEEL_SPEEDS", 0, + {"WHEEL_SPEED_FL": speed, + "WHEEL_SPEED_FR": speed, + "WHEEL_SPEED_RL": speed, + "WHEEL_SPEED_RR": speed}, -1)) + + if engage: + msg.append(packer.make_can_msg("SCM_BUTTONS", 0, {"CRUISE_BUTTONS": 3}, idx)) + else: + msg.append(packer.make_can_msg("SCM_BUTTONS", 0, {"CRUISE_BUTTONS": 0}, idx)) + + values = {"COUNTER_PEDAL": idx&0xF} + checksum = crc8_pedal(packer.make_can_msg("GAS_SENSOR", 0, {"COUNTER_PEDAL": idx&0xF}, -1)[2][:-1]) + values["CHECKSUM_PEDAL"] = checksum + msg.append(packer.make_can_msg("GAS_SENSOR", 0, values, -1)) + + msg.append(packer.make_can_msg("GEARBOX", 0, {"GEAR": 4, "GEAR_SHIFTER": 8}, idx)) + msg.append(packer.make_can_msg("GAS_PEDAL_2", 0, {}, idx)) + msg.append(packer.make_can_msg("SEATBELT_STATUS", 0, {"SEATBELT_DRIVER_LATCHED": 1}, idx)) + msg.append(packer.make_can_msg("STEER_STATUS", 0, {}, idx)) + msg.append(packer.make_can_msg("STEERING_SENSORS", 0, {"STEER_ANGLE": angle}, idx)) + msg.append(packer.make_can_msg("POWERTRAIN_DATA", 0, {}, idx)) + msg.append(packer.make_can_msg("VSA_STATUS", 0, {}, idx)) + msg.append(packer.make_can_msg("STANDSTILL", 0, {}, idx)) + msg.append(packer.make_can_msg("STEER_MOTOR_TORQUE", 0, {}, idx)) + msg.append(packer.make_can_msg("EPB_STATUS", 0, {}, idx)) + msg.append(packer.make_can_msg("DOORS_STATUS", 0, {}, idx)) + msg.append(packer.make_can_msg("CRUISE_PARAMS", 0, {}, idx)) + msg.append(packer.make_can_msg("CRUISE", 0, {}, idx)) + msg.append(packer.make_can_msg("SCM_FEEDBACK", 0, {"MAIN_ON": 1}, idx)) + #print(msg) + + # cam bus + msg.append(packer.make_can_msg("STEERING_CONTROL", 2, {}, idx)) + msg.append(packer.make_can_msg("ACC_HUD", 2, {}, idx)) + msg.append(packer.make_can_msg("BRAKE_COMMAND", 2, {}, idx)) + + # radar + if idx%5 == 0: + msg.append(rpacker.make_can_msg("RADAR_DIAGNOSTIC", 1, {"RADAR_STATE": 0x79}, -1)) + for i in range(16): + msg.append(rpacker.make_can_msg("TRACK_%d" % i, 1, {"LONG_DIST": 255.5}, -1)) + + # fill in the rest for fingerprint + done = set([x[0] for x in msg]) + for k,v in FINGERPRINTS[CAR.CIVIC][0].items(): + if k not in done and k not in [0xE4, 0x194]: + msg.append([k, 0, b'\x00'*v, 0]) + pm.send('can', can_list_to_can_capnp(msg)) + +def sendcan_function(sendcan): + sc = messaging.drain_sock_raw(sendcan) + cp.update_strings(sc, sendcan=True) + + if cp.vl[0x1fa]['COMPUTER_BRAKE_REQUEST']: + brake = cp.vl[0x1fa]['COMPUTER_BRAKE'] * 0.003906248 + else: + brake = 0.0 + + if cp.vl[0x200]['GAS_COMMAND'] > 0: + gas = cp.vl[0x200]['GAS_COMMAND'] / 256.0 + else: + gas = 0.0 + + if cp.vl[0xe4]['STEER_TORQUE_REQUEST']: + steer_torque = cp.vl[0xe4]['STEER_TORQUE']*1.0/0x1000 + else: + steer_torque = 0.0 + + return (gas, brake, steer_torque) + +if __name__ == "__main__": + pm = messaging.PubMaster(['can']) + sendcan = messaging.sub_sock('sendcan') + idx = 0 + while 1: + sendcan_function(sendcan) + can_function(pm, 10.0, idx) + time.sleep(0.01) + idx += 1 + diff --git a/tools/sim/controller.py b/tools/sim/controller.py new file mode 100755 index 0000000000..6afbb1e92f --- /dev/null +++ b/tools/sim/controller.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +import os +import time +import math +import atexit +import numpy as np +import threading +import carla +import random +import cereal.messaging as messaging +from common.params import Params +from common.realtime import Ratekeeper +from can import can_function, sendcan_function +import queue + +pm = messaging.PubMaster(['frame', 'sensorEvents', 'can']) + +W,H = 1164, 874 + +def cam_callback(image): + img = np.frombuffer(image.raw_data, dtype=np.dtype("uint8")) + img = np.reshape(img, (H, W, 4)) + img = img[:, :, [0,1,2]].copy() + + dat = messaging.new_message() + dat.init('frame') + dat.frame = { + "frameId": image.frame, + "image": img.tostring(), + } + pm.send('frame', dat) + +def imu_callback(imu): + #print(imu, imu.accelerometer) + + dat = messaging.new_message() + dat.init('sensorEvents', 2) + dat.sensorEvents[0].sensor = 4 + dat.sensorEvents[0].type = 0x10 + dat.sensorEvents[0].init('acceleration') + dat.sensorEvents[0].acceleration.v = [imu.accelerometer.x, imu.accelerometer.y, imu.accelerometer.z] + # copied these numbers from locationd + dat.sensorEvents[1].sensor = 5 + dat.sensorEvents[1].type = 0x10 + dat.sensorEvents[1].init('gyroUncalibrated') + dat.sensorEvents[1].gyroUncalibrated.v = [imu.gyroscope.x, imu.gyroscope.y, imu.gyroscope.z] + pm.send('sensorEvents', dat) + +def health_function(): + pm = messaging.PubMaster(['health']) + rk = Ratekeeper(1.0) + while 1: + dat = messaging.new_message() + dat.init('health') + dat.valid = True + dat.health = { + 'ignitionLine': True, + 'hwType': "whitePanda", + 'controlsAllowed': True + } + pm.send('health', dat) + rk.keep_time() + +def fake_driver_monitoring(): + pm = messaging.PubMaster(['driverMonitoring']) + while 1: + dat = messaging.new_message() + dat.init('driverMonitoring') + dat.driverMonitoring.faceProb = 1.0 + pm.send('driverMonitoring', dat) + time.sleep(0.1) + +def go(): + client = carla.Client("127.0.0.1", 2000) + client.set_timeout(5.0) + world = client.load_world('Town03') + + settings = world.get_settings() + settings.fixed_delta_seconds = 0.05 + world.apply_settings(settings) + + weather = carla.WeatherParameters( + cloudyness=0.0, + precipitation=0.0, + precipitation_deposits=0.0, + wind_intensity=0.0, + sun_azimuth_angle=0.0, + sun_altitude_angle=0.0) + world.set_weather(weather) + + blueprint_library = world.get_blueprint_library() + """ + for blueprint in blueprint_library.filter('sensor.*'): + print(blueprint.id) + exit(0) + """ + + world_map = world.get_map() + + vehicle_bp = random.choice(blueprint_library.filter('vehicle.bmw.*')) + vehicle = world.spawn_actor(vehicle_bp, random.choice(world_map.get_spawn_points())) + #vehicle.set_autopilot(True) + + blueprint = blueprint_library.find('sensor.camera.rgb') + blueprint.set_attribute('image_size_x', str(W)) + blueprint.set_attribute('image_size_y', str(H)) + blueprint.set_attribute('fov', '70') + blueprint.set_attribute('sensor_tick', '0.05') + transform = carla.Transform(carla.Location(x=0.8, z=1.45)) + camera = world.spawn_actor(blueprint, transform, attach_to=vehicle) + camera.listen(cam_callback) + + # TODO: wait for carla 0.9.7 + imu_bp = blueprint_library.find('sensor.other.imu') + imu = world.spawn_actor(imu_bp, transform, attach_to=vehicle) + imu.listen(imu_callback) + + def destroy(): + print("clean exit") + imu.destroy() + camera.destroy() + vehicle.destroy() + print("done") + atexit.register(destroy) + + threading.Thread(target=health_function).start() + threading.Thread(target=fake_driver_monitoring).start() + + # can loop + sendcan = messaging.sub_sock('sendcan') + rk = Ratekeeper(100) + steer_angle = 0 + while 1: + vel = vehicle.get_velocity() + speed = math.sqrt(vel.x**2 + vel.y**2 + vel.z**2) + + can_function(pm, speed, steer_angle, rk.frame, rk.frame%500 == 499) + if rk.frame%5 == 0: + throttle, brake, steer = sendcan_function(sendcan) + steer_angle += steer/10000.0 # torque + vc = carla.VehicleControl(throttle=throttle, steer=steer_angle, brake=brake) + vehicle.apply_control(vc) + print(speed, steer_angle, vc) + + rk.keep_time() + +if __name__ == "__main__": + params = Params() + params.delete("Offroad_ConnectivityNeeded") + from selfdrive.version import terms_version, training_version + params.put("HasAcceptedTerms", terms_version) + params.put("CompletedTrainingVersion", training_version) + + go() + diff --git a/tools/sim/get.sh b/tools/sim/get.sh new file mode 100755 index 0000000000..187b769a61 --- /dev/null +++ b/tools/sim/get.sh @@ -0,0 +1,10 @@ +#!/bin/bash -e +FILE=CARLA_0.9.7.tar.gz +if [ ! -f $FILE ]; then + curl -O http://carla-assets-internal.s3.amazonaws.com/Releases/Linux/$FILE +fi +mkdir -p carla +cd carla +tar xvf ../$FILE +easy_install PythonAPI/carla/dist/carla-0.9.7-py3.5-linux-x86_64.egg + diff --git a/tools/sim/replay.sh b/tools/sim/replay.sh new file mode 100755 index 0000000000..351d76c9ec --- /dev/null +++ b/tools/sim/replay.sh @@ -0,0 +1,5 @@ +#!/bin/bash +cd ~/one/tools/nui +# vision, boardd, sensorsd, gpsd +ALLOW=frame,can,ubloxRaw,health,sensorEvents,gpsNMEA,gpsLocation ./nui "02ec6bea180a4d36/2019-10-25--10-18-09" + diff --git a/tools/sim/start.sh b/tools/sim/start.sh new file mode 100755 index 0000000000..2d9edfa32e --- /dev/null +++ b/tools/sim/start.sh @@ -0,0 +1,4 @@ +#!/bin/bash +cd carla +./CarlaUE4.sh + diff --git a/tools/ssh/config b/tools/ssh/config new file mode 100644 index 0000000000..a58ed61b47 --- /dev/null +++ b/tools/ssh/config @@ -0,0 +1,9 @@ +Host EON-smays + HostName 192.168.5.11 + Port 8022 + IdentityFile key/id_rsa + +Host EON-wifi + HostName 192.168.43.1 + Port 8022 + IdentityFile key/id_rsa diff --git a/tools/ssh/key/id_rsa b/tools/ssh/key/id_rsa new file mode 100644 index 0000000000..6a8ecfcce9 --- /dev/null +++ b/tools/ssh/key/id_rsa @@ -0,0 +1,28 @@ +-----BEGIN PRIVATE KEY----- +MIIEvAIBADANBgkqhkiG9w0BAQEFAASCBKYwggSiAgEAAoIBAQC+iXXq30Tq+J5N +Kat3KWHCzcmwZ55nGh6WggAqECa5CasBlM9VeROpVu3beA+5h0MibRgbD4DMtVXB +t6gEvZ8nd04E7eLA9LTZyFDZ7SkSOVj4oXOQsT0GnJmKrASW5KslTWqVzTfo2XCt +Z+004ikLxmyFeBO8NOcErW1pa8gFdQDToH9FrA7kgysic/XVESTOoe7XlzRoe/eZ +acEQ+jtnmFd21A4aEADkk00Ahjr0uKaJiLUAPatxs2icIXWpgYtfqqtaKF23wSt6 +1OTu6cAwXbOWr3m+IUSRUO0IRzEIQS3z1jfd1svgzSgSSwZ1Lhj4AoKxIEAIc8qJ +rO4uymCJAgMBAAECggEBAISFevxHGdoL3Z5xkw6oO5SQKO2GxEeVhRzNgmu/HA+q +x8OryqD6O1CWY4037kft6iWxlwiLOdwna2P25ueVM3LxqdQH2KS4DmlCx+kq6FwC +gv063fQPMhC9LpWimvaQSPEC7VUPjQlo4tPY6sTTYBUOh0A1ihRm/x7juKuQCWix +Cq8C/DVnB1X4mGj+W3nJc5TwVJtgJbbiBrq6PWrhvB/3qmkxHRL7dU2SBb2iNRF1 +LLY30dJx/cD73UDKNHrlrsjk3UJc29Mp4/MladKvUkRqNwlYxSuAtJV0nZ3+iFkL +s3adSTHdJpClQer45R51rFDlVsDz2ZBpb/hRNRoGDuECgYEA6A1EixLq7QYOh3cb +Xhyh3W4kpVvA/FPfKH1OMy3ONOD/Y9Oa+M/wthW1wSoRL2n+uuIW5OAhTIvIEivj +6bAZsTT3twrvOrvYu9rx9aln4p8BhyvdjeW4kS7T8FP5ol6LoOt2sTP3T1LOuJPO +uQvOjlKPKIMh3c3RFNWTnGzMPa0CgYEA0jNiPLxP3A2nrX0keKDI+VHuvOY88gdh +0W5BuLMLovOIDk9aQFIbBbMuW1OTjHKv9NK+Lrw+YbCFqOGf1dU/UN5gSyE8lX/Q +FsUGUqUZx574nJZnOIcy3ONOnQLcvHAQToLFAGUd7PWgP3CtHkt9hEv2koUwL4vo +ikTP1u9Gkc0CgYEA2apoWxPZrY963XLKBxNQecYxNbLFaWq67t3rFnKm9E8BAICi +4zUaE5J1tMVi7Vi9iks9Ml9SnNyZRQJKfQ+kaebHXbkyAaPmfv+26rqHKboA0uxA +nDOZVwXX45zBkp6g1sdHxJx8JLoGEnkC9eyvSi0C//tRLx86OhLErXwYcNkCf1it +VMRKrWYoXJTUNo6tRhvodM88UnnIo3u3CALjhgU4uC1RTMHV4ZCGBwiAOb8GozSl +s5YD1E1iKwEULloHnK6BIh6P5v8q7J6uf/xdqoKMjlWBHgq6/roxKvkSPA1DOZ3l +jTadcgKFnRUmc+JT9p/ZbCxkA/ALFg8++G+0ghECgYA8vG3M/utweLvq4RI7l7U7 +b+i2BajfK2OmzNi/xugfeLjY6k2tfQGRuv6ppTjehtji2uvgDWkgjJUgPfZpir3I +RsVMUiFgloWGHETOy0Qvc5AwtqTJFLTD1Wza2uBilSVIEsg6Y83Gickh+ejOmEsY +6co17RFaAZHwGfCFFjO76Q== +-----END PRIVATE KEY----- diff --git a/tools/ssh/via-smays.sh b/tools/ssh/via-smays.sh new file mode 100755 index 0000000000..bf903cb86c --- /dev/null +++ b/tools/ssh/via-smays.sh @@ -0,0 +1,3 @@ +# enp5s0 is the smays network name. Change it appropriately if you are using an ethernet adapter (type ifconfig to get the proper network name) +sudo ifconfig enp5s0 192.168.5.1 netmask 255.255.255.0 +ssh -F config EON-smays diff --git a/tools/ssh/via-wifi.sh b/tools/ssh/via-wifi.sh new file mode 100755 index 0000000000..4a43ca2d83 --- /dev/null +++ b/tools/ssh/via-wifi.sh @@ -0,0 +1 @@ +ssh -F config EON-wifi diff --git a/tools/steer.gif b/tools/steer.gif new file mode 100644 index 0000000000..0c2c96a14a --- /dev/null +++ b/tools/steer.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:34f68c1ee87b92ae6c4f8305afe03c81a12ef3fe328e261639c8b587c65f187c +size 7310573 diff --git a/tools/stream.gif b/tools/stream.gif new file mode 100644 index 0000000000..03947d232e --- /dev/null +++ b/tools/stream.gif @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:b895688a233bbdc8fd8b2915558eb5aa38b66c07d530a74b181a585471a118a7 +size 4605918 diff --git a/tools/streamer/streamerd.py b/tools/streamer/streamerd.py new file mode 100755 index 0000000000..95d34735f1 --- /dev/null +++ b/tools/streamer/streamerd.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python +import os +import sys +import argparse +import zmq +import json +import cv2 +import numpy as np +from hexdump import hexdump +import scipy.misc +import struct +from collections import deque + +# sudo pip install git+git://github.com/mikeboers/PyAV.git +import av + +import cereal.messaging as messaging +from cereal.services import service_list + +PYGAME = os.getenv("PYGAME") is not None +if PYGAME: + import pygame + imgff = np.zeros((874, 1164, 3), dtype=np.uint8) + +# first 74 bytes in any stream +start = "0000000140010c01ffff016000000300b0000003000003005dac5900000001420101016000000300b0000003000003005da0025080381c5c665aee4c92ec80000000014401c0f1800420" + +def receiver_thread(): + if PYGAME: + pygame.init() + pygame.display.set_caption("vnet debug UI") + screen = pygame.display.set_mode((1164,874), pygame.DOUBLEBUF) + camera_surface = pygame.surface.Surface((1164,874), 0, 24).convert() + + addr = "192.168.5.11" + if len(sys.argv) >= 2: + addr = sys.argv[1] + + context = zmq.Context() + s = messaging.sub_sock(context, 9002, addr=addr) + frame_sock = messaging.pub_sock(context, service_list['frame'].port) + + ctx = av.codec.codec.Codec('hevc', 'r').create() + ctx.decode(av.packet.Packet(start.decode("hex"))) + + import time + while 1: + t1 = time.time() + ts, raw = s.recv_multipart() + ts = struct.unpack('q', ts)[0] * 1000 + t1, t2 = time.time(), t1 + #print 'ms to get frame:', (t1-t2)*1000 + + pkt = av.packet.Packet(raw) + f = ctx.decode(pkt) + if not f: + continue + f = f[0] + t1, t2 = time.time(), t1 + #print 'ms to decode:', (t1-t2)*1000 + + y_plane = np.frombuffer(f.planes[0], np.uint8).reshape((874, 1216))[:, 0:1164] + u_plane = np.frombuffer(f.planes[1], np.uint8).reshape((437, 608))[:, 0:582] + v_plane = np.frombuffer(f.planes[2], np.uint8).reshape((437, 608))[:, 0:582] + yuv_img = y_plane.tobytes() + u_plane.tobytes() + v_plane.tobytes() + t1, t2 = time.time(), t1 + #print 'ms to make yuv:', (t1-t2)*1000 + #print 'tsEof:', ts + + dat = messaging.new_message() + dat.init('frame') + dat.frame.image = yuv_img + dat.frame.timestampEof = ts + dat.frame.transform = map(float, list(np.eye(3).flatten())) + frame_sock.send(dat.to_bytes()) + + if PYGAME: + yuv_np = np.frombuffer(yuv_img, dtype=np.uint8).reshape(874 * 3 // 2, -1) + cv2.cvtColor(yuv_np, cv2.COLOR_YUV2RGB_I420, dst=imgff) + #print yuv_np.shape, imgff.shape + + #scipy.misc.imsave("tmp.png", imgff) + + pygame.surfarray.blit_array(camera_surface, imgff.swapaxes(0,1)) + screen.blit(camera_surface, (0, 0)) + pygame.display.flip() + + +def main(gctx=None): + receiver_thread() + +if __name__ == "__main__": + main()