IQ.Lvbs/IQ.Pilot
8
7
Fork 2
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

IQ.Pilot Release Commit @ bec7652

Browse Source
This commit is contained in:
IQ.Lvbs history cleanup
2026-08-22 23:42:41 -05:00
commit 58039e647c
4603 changed files with 1236178 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
iqpilot/tools/lib/__init__.py Normal file
View File

62
iqpilot/tools/lib/api.py Normal file
View File

@@ -0,0 +1,62 @@
import os
import json
import requests
from requests.adapters import HTTPAdapter, Retry
from iqpilot.system.hardware.hw import Paths
API_HOST = os.getenv('API_HOST', 'https://api-iqlabs.konn3kt.com')
# TODO: this should be merged into common.api
class CommaApi:
def __init__(self, token=None):
self.session = requests.Session()
self.session.headers['User-agent'] = 'OpenpilotTools'
if token:
self.session.headers['Authorization'] = 'JWT ' + token
retries = Retry(total=5, backoff_factor=1, status_forcelist=[500, 502, 503, 504])
self.session.mount('https://', HTTPAdapter(max_retries=retries))
def request(self, method, endpoint, **kwargs):
with self.session.request(method, API_HOST + '/' + endpoint, **kwargs) as resp:
resp_json = resp.json()
if isinstance(resp_json, dict) and resp_json.get('error'):
if resp.status_code in [401, 403]:
raise UnauthorizedError('Unauthorized. Authenticate with tools/lib/auth.py')
e = APIError(str(resp.status_code) + ":" + resp_json.get('description', str(resp_json['error'])))
e.status_code = resp.status_code
raise e
return resp_json
def get(self, endpoint, **kwargs):
return self.request('GET', endpoint, **kwargs)
def post(self, endpoint, **kwargs):
return self.request('POST', endpoint, **kwargs)
class APIError(Exception):
pass
class UnauthorizedError(Exception):
pass
def get_token():
try:
with open(os.path.join(Paths.config_root(), 'auth.json')) as f:
return json.load(f)['access_token']
except Exception:
return None
def set_token(token):
os.makedirs(Paths.config_root(), exist_ok=True)
with open(os.path.join(Paths.config_root(), 'auth.json'), 'w') as f:
json.dump({'access_token': token}, f)
def clear_token():
try:
os.unlink(os.path.join(Paths.config_root(), 'auth.json'))
except FileNotFoundError:
pass

114
iqpilot/tools/lib/auth.py Executable file
View File

@@ -0,0 +1,114 @@
#!/usr/bin/env python3
"""
Usage::
usage: auth.py [-h] [{github,jwt}] [jwt]
Login to your konn3kt account
positional arguments:
{github,jwt}
jwt
optional arguments:
-h, --help show this help message and exit
Examples::
./auth.py # Log in with GitHub
./auth.py jwt ey..hw # Log in with a pre-issued JWT (for CI)
"""
import argparse
import sys
import pprint
import webbrowser
from http.server import BaseHTTPRequestHandler, HTTPServer
from typing import Any
from urllib.parse import parse_qs, urlencode
from iqpilot.tools.lib.api import APIError, CommaApi, UnauthorizedError, set_token, get_token
PORT = 3000
class ClientRedirectServer(HTTPServer):
query_params: dict[str, Any] = {}
class ClientRedirectHandler(BaseHTTPRequestHandler):
def do_GET(self):
if '?' in self.path:
query_parsed = parse_qs(self.path.split('?', 1)[1], keep_blank_values=True)
if 'code' in query_parsed or 'error' in query_parsed:
self.server.query_params = query_parsed
self.send_response(200)
self.send_header('Content-type', 'text/plain')
self.end_headers()
self.wfile.write(b'Return to the CLI to continue')
def log_message(self, fmt, *fmt_args):
sys.stderr.write(f"[auth callback] {self.address_string()} {fmt % fmt_args}\n")
def auth_redirect_link(method):
if method != 'github':
raise NotImplementedError(f"no redirect implemented for method {method}")
params = {
'client_id': 'Ov23lifjMafxJzFatvuB',
'redirect_uri': 'https://api-iqlabs.konn3kt.com/v2/auth/h/redirect/',
'state': f'service,localhost:{PORT}',
'scope': 'read:user',
}
return 'https://github.com/login/oauth/authorize?' + urlencode(params)
def login(method):
oauth_uri = auth_redirect_link(method)
web_server = ClientRedirectServer(('localhost', PORT), ClientRedirectHandler)
print(f'To sign in, use your browser and navigate to {oauth_uri}')
webbrowser.open(oauth_uri, new=2)
while True:
web_server.handle_request()
if 'code' in web_server.query_params:
break
elif 'error' in web_server.query_params:
print('Authentication Error: "{}". Description: "{}" '.format(
web_server.query_params['error'],
web_server.query_params.get('error_description')), file=sys.stderr)
break
try:
auth_resp = CommaApi().post('v2/auth/', data={'code': web_server.query_params['code'], 'provider': web_server.query_params['provider']})
set_token(auth_resp['access_token'])
except APIError as e:
print(f'Authentication Error: {e}', file=sys.stderr)
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Login to your konn3kt account')
parser.add_argument('method', default='github', const='github', nargs='?', choices=['github', 'jwt'])
parser.add_argument('jwt', nargs='?')
args = parser.parse_args()
if args.method == 'jwt':
if args.jwt is None:
print("method JWT selected, but no JWT was provided")
exit(1)
set_token(args.jwt)
else:
login(args.method)
try:
me = CommaApi(token=get_token()).get('/v1/me')
print("Authenticated!")
pprint.pprint(me)
except UnauthorizedError:
print("Got invalid JWT")
exit(1)

58
iqpilot/tools/lib/filereader.py Normal file
View File

@@ -0,0 +1,58 @@
import os
import io
import posixpath
import socket
from functools import cache
from iqpilot.common.utils import retry
from urllib.parse import urlparse
from iqpilot.tools.lib.url_file import URLFile
DATA_ENDPOINT = os.getenv("DATA_ENDPOINT", "http://data-raw.comma.internal/")
@cache
@retry(delay=0.0)
def internal_source_available(url: str) -> bool:
if os.path.isdir(url):
return True
try:
hostname = urlparse(url).hostname
port = urlparse(url).port or 80
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.settimeout(0.5)
s.connect((hostname, port))
return True
except (socket.gaierror, ConnectionRefusedError):
pass
return False
def resolve_name(fn):
if fn.startswith("cd:/"):
return posixpath.join(DATA_ENDPOINT, fn[4:])
return fn
@cache
def file_exists(fn):
fn = resolve_name(fn)
if fn.startswith(("http://", "https://")):
return URLFile(fn).get_length_online() != -1
return os.path.exists(fn)
class DiskFile(io.BufferedReader):
def get_multi_range(self, ranges: list[tuple[int, int]]) -> list[bytes]:
parts = []
for r in ranges:
self.seek(r[0])
parts.append(self.read(r[1] - r[0]))
return parts
def FileReader(fn):
fn = resolve_name(fn)
if fn.startswith(("http://", "https://")):
return URLFile(fn)
else:
return DiskFile(open(fn, "rb"))

176
iqpilot/tools/lib/framereader.py Normal file
View File

@@ -0,0 +1,176 @@
import os
import subprocess
import json
import logging
from collections.abc import Iterator
from collections import OrderedDict
import numpy as np
from iqpilot.tools.lib.filereader import FileReader, resolve_name
from iqpilot.tools.lib.vidindex import hevc_index
class DataUnreadableError(Exception):
pass
logger = logging.getLogger("tools")
HEVC_SLICE_B = 0
HEVC_SLICE_P = 1
HEVC_SLICE_I = 2
class LRUCache:
def __init__(self, capacity: int):
self._cache: OrderedDict = OrderedDict()
self.capacity = capacity
def __getitem__(self, key):
self._cache.move_to_end(key)
return self._cache[key]
def __setitem__(self, key, value):
self._cache[key] = value
if len(self._cache) > self.capacity:
self._cache.popitem(last=False)
def __contains__(self, key):
return key in self._cache
def assert_hvec(fn: str) -> None:
with FileReader(fn) as f:
header = f.read(4)
if len(header) == 0:
raise DataUnreadableError(f"{fn} is empty")
elif header == b"\x00\x00\x00\x01":
if 'hevc' not in fn:
raise NotImplementedError(fn)
def decompress_video_data(rawdat, w, h, pix_fmt="rgb24", vid_fmt='hevc', hwaccel="auto", loglevel="info") -> np.ndarray:
threads = os.getenv("FFMPEG_THREADS", "0")
args = ["ffmpeg", "-v", loglevel,
"-threads", threads,
"-hwaccel", hwaccel,
"-c:v", "hevc",
"-vsync", "0",
"-f", vid_fmt,
"-flags2", "showall",
"-i", "-",
"-f", "rawvideo",
"-pix_fmt", pix_fmt,
"-"]
dat = subprocess.check_output(args, input=rawdat)
ret: np.ndarray
if pix_fmt == "rgb24":
ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, h, w, 3)
elif pix_fmt in ["nv12", "yuv420p"]:
ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, (h*w*3//2))
else:
raise NotImplementedError(f"Unsupported pixel format: {pix_fmt}")
return ret
def ffprobe(fn, fmt=None):
fn = resolve_name(fn)
cmd = ["ffprobe", "-v", "quiet", "-print_format", "json", "-show_format", "-show_streams"]
if fmt:
cmd += ["-f", fmt]
cmd += ["-i", "-"]
try:
with FileReader(fn) as f:
ffprobe_output = subprocess.check_output(cmd, input=f.read(4096))
except subprocess.CalledProcessError as e:
raise DataUnreadableError(fn) from e
return json.loads(ffprobe_output)
def get_index_data(fn: str, index_data: dict|None = None):
if index_data is None:
index_data = get_video_index(fn)
if index_data is None:
raise DataUnreadableError(f"Failed to index {fn!r}")
stream = index_data["probe"]["streams"][0]
return index_data["index"], index_data["global_prefix"], stream["width"], stream["height"]
def get_video_index(fn):
assert_hvec(fn)
frame_types, dat_len, prefix = hevc_index(fn)
index = np.array(frame_types + [(0xFFFFFFFF, dat_len)], dtype=np.uint32)
probe = ffprobe(fn, "hevc")
return {
'index': index,
'global_prefix': prefix,
'probe': probe
}
class FfmpegDecoder:
def __init__(self, fn: str, index_data: dict|None = None,
pix_fmt: str = "rgb24", hwaccel="auto", loglevel="quiet"):
self.fn = fn
self.index, self.prefix, self.w, self.h = get_index_data(fn, index_data)
self.frame_count = len(self.index) - 1 # sentinel row at the end
self.iframes = np.where(self.index[:, 0] == HEVC_SLICE_I)[0]
self.pix_fmt = pix_fmt
self.loglevel, self.hwaccel = loglevel, hwaccel
def _gop_bounds(self, frame_idx: int):
f_b = frame_idx
while f_b > 0 and self.index[f_b, 0] != HEVC_SLICE_I:
f_b -= 1
f_e = frame_idx + 1
while f_e < self.frame_count and self.index[f_e, 0] != HEVC_SLICE_I:
f_e += 1
return f_b, f_e, self.index[f_b, 1], self.index[f_e, 1]
def _decode_gop(self, raw: bytes) -> Iterator[np.ndarray]:
yield from decompress_video_data(raw, self.w, self.h, pix_fmt=self.pix_fmt, hwaccel=self.hwaccel, loglevel=self.loglevel)
def get_gop_start(self, frame_idx: int):
return self.iframes[np.searchsorted(self.iframes, frame_idx, side="right") - 1]
def get_iterator(self, start_fidx: int = 0, end_fidx: int|None = None,
frame_skip: int = 1) -> Iterator[tuple[int, np.ndarray]]:
end_fidx = end_fidx or self.frame_count
fidx = start_fidx
while fidx < end_fidx:
f_b, f_e, off_b, off_e = self._gop_bounds(fidx)
with FileReader(self.fn) as f:
f.seek(off_b)
raw = self.prefix + f.read(off_e - off_b)
# number of frames to discard inside this GOP before the wanted one
for i, frm in enumerate(decompress_video_data(raw, self.w, self.h, self.pix_fmt, hwaccel=self.hwaccel, loglevel=self.loglevel)):
fidx = f_b + i
if fidx >= end_fidx:
return
elif fidx >= start_fidx and (fidx - start_fidx) % frame_skip == 0:
yield fidx, frm
fidx += 1
def FrameIterator(fn: str, index_data: dict|None=None, pix_fmt: str = "rgb24",
start_fidx:int=0, end_fidx=None, frame_skip:int=1, hwaccel="auto", loglevel="quiet") -> Iterator[np.ndarray]:
dec = FfmpegDecoder(fn, pix_fmt=pix_fmt, index_data=index_data, hwaccel=hwaccel, loglevel=loglevel)
for _, frame in dec.get_iterator(start_fidx=start_fidx, end_fidx=end_fidx, frame_skip=frame_skip):
yield frame
class FrameReader:
def __init__(self, fn: str, index_data: dict|None = None, cache_size: int = 30,
pix_fmt: str = "rgb24", hwaccel="auto", loglevel="quiet"):
self.decoder = FfmpegDecoder(fn, index_data=index_data, pix_fmt=pix_fmt, hwaccel=hwaccel, loglevel=loglevel)
self.iframes = self.decoder.iframes
self._cache: LRUCache = LRUCache(cache_size)
self.w, self.h, self.frame_count, = self.decoder.w, self.decoder.h, self.decoder.frame_count
self.pix_fmt = pix_fmt
self.it: Iterator[tuple[int, np.ndarray]] | None = None
self.fidx = -1
def get(self, fidx:int):
if fidx in self._cache: # If frame is cached, return it
return self._cache[fidx]
read_start = self.decoder.get_gop_start(fidx)
if not self.it or fidx < self.fidx or read_start != self.decoder.get_gop_start(self.fidx): # If the frame is in a different GOP, reset the iterator
self.it = self.decoder.get_iterator(read_start)
self.fidx = -1
while self.fidx < fidx:
self.fidx, frame = next(self.it)
self._cache[self.fidx] = frame
return self._cache[fidx]

436
iqpilot/tools/lib/logreader.py Executable file
View File

@@ -0,0 +1,436 @@
#!/usr/bin/env python3
import bz2
from functools import partial
import multiprocessing
import capnp
import enum
import os
import pathlib
import sys
import tqdm
import urllib.parse
import warnings
import zstandard as zstd
import numpy as np
from collections.abc import Callable, Iterable, Iterator
from pathlib import Path
from urllib.parse import parse_qs, urlparse
from iqpilot.cereal import log as capnp_log, messaging
from iqpilot.cereal.services import SERVICE_LIST
from iqpilot.common.swaglog import cloudlog
from iqpilot.tools.lib.filereader import DATA_ENDPOINT, FileReader, file_exists, internal_source_available
from iqpilot.tools.lib.route import Route, SegmentRange, FileName
LogMessage = type[capnp._DynamicStructReader]
LogIterable = Iterable[LogMessage]
RawLogIterable = Iterable[bytes]
FileNames = tuple[str, ...]
Source = Callable[[SegmentRange, list[int], FileNames], dict[int, str]]
InternalUnavailableException = Exception("Internal source not available")
OPENPILOT_CI_BASE_URL = "https://commadataci.blob.core.windows.net/openpilotci/"
OPENPILOT_CI_ACCOUNT_URL = "https://commadataci.blob.core.windows.net"
def get_url(route_name: str, segment_num: str | int, filename: str) -> str:
return f"{OPENPILOT_CI_BASE_URL}{route_name.replace('|', '/')}/{segment_num}/{filename}"
def upload_file(path: str, blob_name: str, overwrite=False) -> str:
from azure.identity import AzureCliCredential
from azure.storage.blob import BlobClient
token_path = Path("/data/azure_token")
credential = os.environ.get("AZURE_TOKEN") or (token_path.read_text().strip() if token_path.is_file() else AzureCliCredential())
client = BlobClient(OPENPILOT_CI_ACCOUNT_URL, container_name="openpilotci", blob_name=blob_name, credential=credential)
with open(path, "rb") as f:
client.upload_blob(f, overwrite=overwrite)
return OPENPILOT_CI_BASE_URL + blob_name
def comma_api_source(sr: SegmentRange, seg_idxs: list[int], fns: FileNames) -> dict[int, str]:
route = Route(sr.route_name)
if fns == FileName.RLOG:
return {seg: route.log_paths()[seg] for seg in seg_idxs if route.log_paths()[seg] is not None}
return {seg: route.qlog_paths()[seg] for seg in seg_idxs if route.qlog_paths()[seg] is not None}
def internal_source(sr: SegmentRange, seg_idxs: list[int], fns: FileNames, endpoint_url: str = DATA_ENDPOINT) -> dict[int, str]:
if not internal_source_available(endpoint_url):
raise InternalUnavailableException
def internal_url(seg, file):
return f"{endpoint_url.rstrip('/')}/{sr.dongle_id}/{sr.log_id}/{seg}/{file}"
return eval_source({seg: [internal_url(seg, fn) for fn in fns] for seg in seg_idxs})
def openpilotci_source(sr: SegmentRange, seg_idxs: list[int], fns: FileNames) -> dict[int, str]:
return eval_source({seg: [get_url(sr.route_name, seg, fn) for fn in fns] for seg in seg_idxs})
def eval_source(files: dict[int, list[str] | str]) -> dict[int, str]:
valid_files: dict[int, str] = {}
for seg_idx, urls in files.items():
if isinstance(urls, str):
urls = [urls]
for url in urls:
if file_exists(url):
valid_files[seg_idx] = url
break
return valid_files
ALL_SERVICES = list(SERVICE_LIST.keys())
def raw_live_logreader(services: list[str] = ALL_SERVICES, addr: str = '127.0.0.1') -> RawLogIterable:
if addr != "127.0.0.1":
os.environ["ZMQ"] = "1"
messaging.reset_context()
poller = messaging.Poller()
for service in services:
messaging.sub_sock(service, poller, addr=addr)
while True:
for sock in poller.poll(100):
yield sock.receive()
def live_logreader(services: list[str] = ALL_SERVICES, addr: str = '127.0.0.1') -> LogIterable:
for msg in raw_live_logreader(services, addr):
with capnp_log.Event.from_bytes(msg) as evt:
yield evt
def flatten_type_dict(data, sep="/", prefix=None):
result = {}
if isinstance(data, dict):
for key, value in data.items():
result.update(flatten_type_dict(value, sep, key if prefix is None else prefix + sep + key))
return result
if isinstance(data, list):
return {prefix: np.array(data)}
return {prefix: data}
def get_message_dict(message, typ):
valid = message.valid
message = message._get(typ)
if not hasattr(message, 'to_dict') or typ in ('qcomGnss', 'ubloxGnss'):
return None
result = flatten_type_dict(message.to_dict(verbose=True))
result['_valid'] = valid
return result
def potentially_ragged_array(values, dtype=None, **kwargs):
try:
return np.array(values, dtype=dtype, **kwargs)
except ValueError:
return np.array(values, dtype=object, **kwargs)
def msgs_to_time_series(msgs):
values = {}
for msg in msgs:
typ = msg.which()
msg_dict = get_message_dict(msg, typ)
if msg_dict is None:
continue
group = values.setdefault(typ, {"t": [], **{key: [] for key in msg_dict}})
group["t"].append(msg.logMonoTime / 1.0e9)
for key, value in msg_dict.items():
group[key].append(value)
for group in values.values():
order = np.argsort(group["t"])
for name, group_values in group.items():
group[name] = potentially_ragged_array(group_values)[order]
return values
def save_log(dest, log_msgs, compress=True):
dat = b"".join(msg.as_builder().to_bytes() for msg in log_msgs)
if compress and dest.endswith(".bz2"):
dat = bz2.compress(dat)
elif compress and dest.endswith(".zst"):
dat = zstd.compress(dat, 10)
with open(dest, "wb") as f:
f.write(dat)
def decompress_stream(data: bytes):
dctx = zstd.ZstdDecompressor()
decompressed_data = b""
with dctx.stream_reader(data) as reader:
decompressed_data = reader.read()
return decompressed_data
class CachedEventReader:
__slots__ = ('_evt', '_enum')
def __init__(self, evt: capnp._DynamicStructReader, _enum: str | None = None):
"""All capnp attribute accesses are expensive, and which() is often called multiple times"""
self._evt = evt
self._enum: str | None = _enum
# fast pickle support
def __reduce__(self):
return CachedEventReader._reducer, (self._evt.as_builder().to_bytes(), self._enum)
@staticmethod
def _reducer(data: bytes, _enum: str | None = None):
with capnp_log.Event.from_bytes(data) as evt:
return CachedEventReader(evt, _enum)
def __repr__(self):
return self._evt.__repr__()
def __str__(self):
return self._evt.__str__()
def __dir__(self):
return dir(self._evt)
def which(self) -> str:
if self._enum is None:
self._enum = self._evt.which()
return self._enum
def __getattr__(self, name: str):
if name.startswith("__") and name.endswith("__"):
return getattr(self, name)
return getattr(self._evt, name)
class _LogFileReader:
def __init__(self, fn, only_union_types=False, sort_by_time=False, dat=None):
self.data_version = None
self._only_union_types = only_union_types
ext = None
if not dat:
_, ext = os.path.splitext(urllib.parse.urlparse(fn).path)
if ext not in ('', '.bz2', '.zst'):
# old rlogs weren't compressed
raise ValueError(f"unknown extension {ext}")
with FileReader(fn) as f:
dat = f.read()
if ext == ".bz2" or dat.startswith(b'BZh9'):
dat = bz2.decompress(dat)
elif ext == ".zst" or dat.startswith(b'\x28\xB5\x2F\xFD'):
# https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#zstandard-frames
dat = decompress_stream(dat)
ents = capnp_log.Event.read_multiple_bytes(dat)
self._ents = []
try:
for e in ents:
self._ents.append(CachedEventReader(e))
except capnp.KjException:
warnings.warn("Corrupted events detected", RuntimeWarning, stacklevel=1)
if sort_by_time:
self._ents.sort(key=lambda x: x.logMonoTime)
def __iter__(self) -> Iterator[capnp._DynamicStructReader]:
for ent in self._ents:
if self._only_union_types:
try:
ent.which()
yield ent
except (capnp.lib.capnp.KjException, RuntimeError):
pass
else:
yield ent
class ReadMode(enum.StrEnum):
RLOG = "r" # only read rlogs
QLOG = "q" # only read qlogs
AUTO = "a" # default to rlogs, fallback to qlogs
AUTO_INTERACTIVE = "i" # default to rlogs, fallback to qlogs with a prompt from the user
class LogsUnavailable(Exception):
pass
def direct_source(file_or_url: str) -> list[str]:
return [file_or_url]
# TODO this should apply to camera files as well
def auto_source(identifier: str, sources: list[Source], default_mode: ReadMode) -> list[str]:
exceptions = {}
sr = SegmentRange(identifier)
needed_seg_idxs = sr.seg_idxs
mode = default_mode if sr.selector is None else ReadMode(sr.selector)
if mode == ReadMode.QLOG:
try_fns = [FileName.QLOG]
else:
try_fns = [FileName.RLOG]
# If selector allows it, fallback to qlogs
if mode in (ReadMode.AUTO, ReadMode.AUTO_INTERACTIVE):
try_fns.append(FileName.QLOG)
# Build a dict of valid files as we evaluate each source. May contain mix of rlogs, qlogs, and None.
# This function only returns when we've sourced all files, or throws an exception
valid_files: dict[int, str] = {}
for fn in try_fns:
for source in sources:
try:
files = source(sr, needed_seg_idxs, fn)
# Build a dict of valid files
valid_files |= files
# Don't check for segment files that have already been found
needed_seg_idxs = [idx for idx in needed_seg_idxs if idx not in valid_files]
# We've found all files, return them
if len(needed_seg_idxs) == 0:
return list(valid_files.values())
else:
raise FileNotFoundError(f"Did not find {fn} for seg idxs {needed_seg_idxs} of {sr.route_name}")
except Exception as e:
exceptions[source.__name__] = e
if fn == try_fns[0]:
missing_logs = len(needed_seg_idxs)
if mode == ReadMode.AUTO:
cloudlog.warning(f"{missing_logs}/{len(sr.seg_idxs)} rlogs were not found, falling back to qlogs for those segments...")
elif mode == ReadMode.AUTO_INTERACTIVE:
if input(f"{missing_logs}/{len(sr.seg_idxs)} rlogs were not found, would you like to fallback to qlogs for those segments? (y/N) ").lower() != "y":
break
missing_logs = len(needed_seg_idxs)
raise LogsUnavailable(f"{missing_logs}/{len(sr.seg_idxs)} logs were not found, please ensure all logs " +
"are uploaded. You can fall back to qlogs with '/a' selector at the end of the route name.\n\n" +
"Exceptions for sources:\n - " + "\n - ".join([f"{k}: {repr(v)}" for k, v in exceptions.items()]))
def parse_indirect(identifier: str) -> str:
if "useradmin.comma.ai" in identifier:
query = parse_qs(urlparse(identifier).query)
identifier = query["onebox"][0]
elif "connect.comma.ai" in identifier or "konn3kt.com" in identifier:
path = urlparse(identifier).path.strip("/").split("/")
if path and path[0] == "connectdata":
# signed data URL from the API host (api-*.konn3kt.com/connectdata/...), not a share link
return identifier
path = ['/'.join(path[:2]), *path[2:]] # recombine log id
identifier = path[0]
if len(path) > 2:
# convert url with seconds to segments
start, end = int(path[1]) // 60, int(path[2]) // 60 + 1
identifier = f"{identifier}/{start}:{end}"
# add selector if it exists
if len(path) > 3:
identifier += f"/{path[3]}"
else:
# add selector if it exists
identifier = "/".join(path)
return identifier
def parse_direct(identifier: str):
if identifier.startswith(("http://", "https://", "cd:/")) or pathlib.Path(identifier).exists():
return identifier
return None
class LogReader:
def _parse_identifier(self, identifier: str) -> list[str]:
# useradmin, etc.
identifier = parse_indirect(identifier)
# direct url or file
direct_parsed = parse_direct(identifier)
if direct_parsed is not None:
return direct_source(identifier)
identifiers = auto_source(identifier, self.sources, self.default_mode)
return identifiers
def __init__(self, identifier: str | list[str], default_mode: ReadMode = ReadMode.RLOG,
sources: list[Source] | None = None, sort_by_time=False, only_union_types=False):
if sources is None:
sources = [internal_source, comma_api_source, openpilotci_source]
self.default_mode = default_mode
self.sources = sources
self.identifier = identifier
if isinstance(identifier, str):
self.identifier = [identifier]
self.sort_by_time = sort_by_time
self.only_union_types = only_union_types
self.__lrs: dict[int, _LogFileReader] = {}
self.reset()
def _get_lr(self, i):
if i not in self.__lrs:
self.__lrs[i] = _LogFileReader(self.logreader_identifiers[i], sort_by_time=self.sort_by_time, only_union_types=self.only_union_types)
return self.__lrs[i]
def __iter__(self):
for i in range(len(self.logreader_identifiers)):
yield from self._get_lr(i)
def _run_on_segment(self, func, i):
return func(self._get_lr(i))
def run_across_segments(self, num_processes, func, disable_tqdm=False, desc=None):
with multiprocessing.Pool(num_processes) as pool:
ret = []
num_segs = len(self.logreader_identifiers)
for p in tqdm.tqdm(pool.imap(partial(self._run_on_segment, func), range(num_segs)), total=num_segs, disable=disable_tqdm, desc=desc):
ret.extend(p)
return ret
def reset(self):
self.logreader_identifiers = []
for identifier in self.identifier:
self.logreader_identifiers.extend(self._parse_identifier(identifier))
@staticmethod
def from_bytes(dat):
return _LogFileReader("", dat=dat)
def filter(self, msg_type: str):
return (getattr(m, m.which()) for m in filter(lambda m: m.which() == msg_type, self))
def first(self, msg_type: str):
return next(self.filter(msg_type), None)
@property
def time_series(self):
return msgs_to_time_series(self)
if __name__ == "__main__":
import codecs
# capnproto <= 0.8.0 throws errors converting byte data to string
# below line catches those errors and replaces the bytes with \x__
codecs.register_error("strict", codecs.backslashreplace_errors)
log_path = sys.argv[1]
lr = LogReader(log_path, sort_by_time=True)
for msg in lr:
print(msg)

382
iqpilot/tools/lib/route.py Normal file
View File

@@ -0,0 +1,382 @@
import os
import re
import requests
from functools import cache
from urllib.parse import urlparse
from collections import defaultdict
from itertools import chain
from iqpilot.tools.lib.api import APIError, CommaApi, get_token
class RE:
DONGLE_ID = r'(?P<dongle_id>[a-f0-9]{16})'
TIMESTAMP = r'(?P<timestamp>[0-9]{4}-[0-9]{2}-[0-9]{2}--[0-9]{2}-[0-9]{2}-[0-9]{2})'
LOG_ID_V2 = r'(?P<count>[a-f0-9]{8})--(?P<uid>[a-z0-9]{10})'
LOG_ID = fr'(?P<log_id>(?:{TIMESTAMP}|{LOG_ID_V2}))'
ROUTE_NAME = fr'(?P<route_name>{DONGLE_ID}[|_/]{LOG_ID})'
SEGMENT_NAME = fr'{ROUTE_NAME}(?:--|/)(?P<segment_num>[0-9]+)'
INDEX = r'-?[0-9]+'
SLICE = fr'(?P<start>{INDEX})?:?(?P<end>{INDEX})?:?(?P<step>{INDEX})?'
SEGMENT_RANGE = fr'{ROUTE_NAME}(?:(--|/)(?P<slice>({SLICE})))?(?:/(?P<selector>([qra])))?'
BOOTLOG_NAME = ROUTE_NAME
EXPLORER_FILE = fr'^(?P<segment_name>{SEGMENT_NAME})--(?P<file_name>[a-z]+\.[a-z0-9]+)$'
OP_SEGMENT_DIR = fr'^(?P<segment_name>{SEGMENT_NAME})$'
class FileName:
RLOG = ("rlog.zst", "rlog.bz2")
QLOG = ("qlog.zst", "qlog.bz2")
QCAMERA = ('qcamera.ts',)
FCAMERA = ('fcamera.hevc',)
ECAMERA = ('ecamera.hevc',)
DCAMERA = ('dcamera.hevc',)
BOOTLOG = ('bootlog.zst', 'bootlog.bz2')
class Route:
def __init__(self, name, data_dir=None):
self._name = RouteName(name)
self.files = None
if data_dir is not None:
self._segments = self._get_segments_local(data_dir)
else:
self._segments = self._get_segments_remote()
self.max_seg_number = self._segments[-1].name.segment_num
@property
def name(self):
return self._name
@property
def segments(self):
return self._segments
def log_paths(self):
log_path_by_seg_num = {s.name.segment_num: s.log_path for s in self._segments}
return [log_path_by_seg_num.get(i, None) for i in range(self.max_seg_number + 1)]
def qlog_paths(self):
qlog_path_by_seg_num = {s.name.segment_num: s.qlog_path for s in self._segments}
return [qlog_path_by_seg_num.get(i, None) for i in range(self.max_seg_number + 1)]
def camera_paths(self):
camera_path_by_seg_num = {s.name.segment_num: s.camera_path for s in self._segments}
return [camera_path_by_seg_num.get(i, None) for i in range(self.max_seg_number + 1)]
def dcamera_paths(self):
dcamera_path_by_seg_num = {s.name.segment_num: s.dcamera_path for s in self._segments}
return [dcamera_path_by_seg_num.get(i, None) for i in range(self.max_seg_number + 1)]
def ecamera_paths(self):
ecamera_path_by_seg_num = {s.name.segment_num: s.ecamera_path for s in self._segments}
return [ecamera_path_by_seg_num.get(i, None) for i in range(self.max_seg_number + 1)]
def qcamera_paths(self):
qcamera_path_by_seg_num = {s.name.segment_num: s.qcamera_path for s in self._segments}
return [qcamera_path_by_seg_num.get(i, None) for i in range(self.max_seg_number + 1)]
# TODO: refactor this, it's super repetitive
def _get_segments_remote(self):
api = CommaApi(get_token())
route_files = api.get('v1/route/' + self.name.canonical_name + '/files')
self.files = [f['url'] if isinstance(f, dict) else f
for f in chain.from_iterable(route_files.values())]
segments = {}
for url in self.files:
_, dongle_id, time_str, segment_num, fn = urlparse(url).path.rsplit('/', maxsplit=4)
segment_name = f'{dongle_id}|{time_str}--{segment_num}'
if segments.get(segment_name):
segments[segment_name] = Segment(
segment_name,
url if fn in FileName.RLOG else segments[segment_name].log_path,
url if fn in FileName.QLOG else segments[segment_name].qlog_path,
url if fn in FileName.FCAMERA else segments[segment_name].camera_path,
url if fn in FileName.DCAMERA else segments[segment_name].dcamera_path,
url if fn in FileName.ECAMERA else segments[segment_name].ecamera_path,
url if fn in FileName.QCAMERA else segments[segment_name].qcamera_path,
)
else:
segments[segment_name] = Segment(
segment_name,
url if fn in FileName.RLOG else None,
url if fn in FileName.QLOG else None,
url if fn in FileName.FCAMERA else None,
url if fn in FileName.DCAMERA else None,
url if fn in FileName.ECAMERA else None,
url if fn in FileName.QCAMERA else None,
)
return sorted(segments.values(), key=lambda seg: seg.name.segment_num)
def _get_segments_local(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(RE.EXPLORER_FILE, f)
op_match = re.match(RE.OP_SEGMENT_DIR, f)
if explorer_match:
segment_name = explorer_match.group('segment_name')
fn = explorer_match.group('file_name')
if segment_name.replace('_', '|').startswith(self.name.canonical_name):
segment_files[segment_name].append((fullpath, fn))
elif op_match and os.path.isdir(fullpath):
segment_name = op_match.group('segment_name')
if segment_name.startswith(self.name.canonical_name):
for seg_f in os.listdir(fullpath):
segment_files[segment_name].append((os.path.join(fullpath, seg_f), seg_f))
elif f == self.name.canonical_name:
for seg_num in os.listdir(fullpath):
if not seg_num.isdigit():
continue
segment_name = f'{self.name.canonical_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 FileName.RLOG)
except StopIteration:
log_path = None
try:
qlog_path = next(path for path, filename in files if filename in FileName.QLOG)
except StopIteration:
qlog_path = None
try:
camera_path = next(path for path, filename in files if filename in FileName.FCAMERA)
except StopIteration:
camera_path = None
try:
dcamera_path = next(path for path, filename in files if filename in FileName.DCAMERA)
except StopIteration:
dcamera_path = None
try:
ecamera_path = next(path for path, filename in files if filename in FileName.ECAMERA)
except StopIteration:
ecamera_path = None
try:
qcamera_path = next(path for path, filename in files if filename in FileName.QCAMERA)
except StopIteration:
qcamera_path = None
segments.append(Segment(segment, log_path, qlog_path, camera_path, dcamera_path, ecamera_path, qcamera_path))
if len(segments) == 0:
raise ValueError(f'Could not find segments for route {self.name.canonical_name} in data directory {data_dir}')
return sorted(segments, key=lambda seg: seg.name.segment_num)
class Segment:
def __init__(self, name, log_path, qlog_path, camera_path, dcamera_path, ecamera_path, qcamera_path):
self._events = None
self._name = SegmentName(name)
self.log_path = log_path
self.qlog_path = qlog_path
self.camera_path = camera_path
self.dcamera_path = dcamera_path
self.ecamera_path = ecamera_path
self.qcamera_path = qcamera_path
@property
def name(self):
return self._name
@staticmethod
@cache
def _get_route_metadata(route_name: str):
api = CommaApi(get_token())
return api.get(f'v1/route/{route_name}')
@property
def url(self):
route_name = self._name.route_name.canonical_name
metadata = self._get_route_metadata(route_name)
return f'{metadata["url"]}/{self._name.segment_num}'
@property
def events(self):
if not self._events:
try:
resp = requests.get(f'{self.url}/events.json')
resp.raise_for_status()
self._events = resp.json()
except Exception as e:
raise APIError(f'error getting events for segment {self._name}') from e
return self._events
class RouteName:
def __init__(self, name_str: str):
self._name_str = name_str
delim = next(c for c in self._name_str if c in ("|", "/"))
self._dongle_id, self._time_str = self._name_str.split(delim)
assert len(self._dongle_id) == 16, self._name_str
assert len(self._time_str) == 20, self._name_str
self._canonical_name = f"{self._dongle_id}|{self._time_str}"
@property
def canonical_name(self) -> str: return self._canonical_name
@property
def dongle_id(self) -> str: return self._dongle_id
@property
def log_id(self) -> str: return self._time_str
@property
def time_str(self) -> str: return self._time_str
@property
def azure_prefix(self):
return f'{self.dongle_id}/{self.log_id}'
def __str__(self) -> str: return self._canonical_name
class SegmentName:
# TODO: add constructor that takes dongle_id, time_str, segment_num and then create instances
# of this class instead of manually constructing a segment name (use canonical_name prop instead)
def __init__(self, name_str: str, allow_route_name=False):
data_dir_path_separator_index = name_str.rsplit("|", 1)[0].rfind("/")
use_data_dir = (data_dir_path_separator_index != -1) and ("|" in name_str)
self._name_str = name_str[data_dir_path_separator_index + 1:] if use_data_dir else name_str
self._data_dir = name_str[:data_dir_path_separator_index] if use_data_dir else None
seg_num_delim = "--" if self._name_str.count("--") == 2 else "/"
name_parts = self._name_str.rsplit(seg_num_delim, 1)
if allow_route_name and len(name_parts) == 1:
name_parts.append("-1") # no segment number
self._route_name = RouteName(name_parts[0])
self._num = int(name_parts[1])
self._canonical_name = f"{self._route_name._dongle_id}|{self._route_name._time_str}--{self._num}"
@property
def canonical_name(self) -> str: return self._canonical_name
# TODO should only use one name
@property
def data_name(self) -> str: return f"{self._route_name.canonical_name}/{self._num}"
@property
def azure_prefix(self):
return f'{self.dongle_id}/{self.log_id}/{self._num}'
@property
def dongle_id(self) -> str: return self._route_name.dongle_id
@property
def time_str(self) -> str: return self._route_name.time_str
@property
def log_id(self) -> str: return self._route_name.time_str
@property
def segment_num(self) -> int: return self._num
@property
def route_name(self) -> RouteName: return self._route_name
@property
def data_dir(self) -> str | None: return self._data_dir
def __str__(self) -> str: return self._canonical_name
@staticmethod
def from_file_name(file_name):
# ??????/xxxxxxxxxxxxxxxx|1111-11-11-11--11-11-11/1/rlog.bz2
dongle_id, route_name, segment_num = file_name.replace('|', '/').split('/')[-4:-1]
return SegmentName(dongle_id + "|" + route_name + "--" + segment_num)
@staticmethod
def from_device_key(dongle_id, key):
# 2018-05-07--18-56-13--5/rlog.bz2
segment_name = key.split('/')[0]
return SegmentName(dongle_id + "|" + segment_name)
@staticmethod
def from_file_key(key):
# 38c52c217150700f/2018-05-07--18-56-13/5/rlog.bz2
az_prefix = '/'.join(key.split('/')[:3])
return SegmentName.from_azure_prefix(az_prefix)
@staticmethod
def from_azure_prefix(prefix):
# xxxxxxxx/1111-11-11-11--11-11-11/0
dongle_id, route_name, segment_num = prefix.split("/")
return SegmentName(dongle_id + "|" + route_name + "--" + segment_num)
@cache
def get_max_seg_number_cached(sr: 'SegmentRange') -> int:
try:
api = CommaApi(get_token())
max_seg_number = api.get("/v1/route/" + sr.route_name.replace("/", "|"))["maxqlog"]
assert isinstance(max_seg_number, int)
return max_seg_number
except Exception as e:
raise Exception("unable to get max_segment_number. ensure you have access to this route or the route is public.") from e
class SegmentRange:
def __init__(self, segment_range: str):
m = re.fullmatch(RE.SEGMENT_RANGE, segment_range)
assert m is not None, f"Segment range is not valid {segment_range}"
self.m = m
@property
def route_name(self) -> str:
return self.m.group("route_name")
@property
def dongle_id(self) -> str:
return self.m.group("dongle_id")
@property
def log_id(self) -> str:
return self.m.group("log_id")
@property
def slice(self) -> str:
return self.m.group("slice") or ""
@property
def selector(self) -> str | None:
return self.m.group("selector")
@property
def seg_idxs(self) -> list[int]:
m = re.fullmatch(RE.SLICE, self.slice)
assert m is not None, f"Invalid slice: {self.slice}"
start, end, step = (None if s is None else int(s) for s in m.groups())
# one segment specified
if start is not None and end is None and ':' not in self.slice:
if start < 0:
start += get_max_seg_number_cached(self) + 1
return [start]
s = slice(start, end, step)
# no specified end or using relative indexing, need number of segments
if end is None or end < 0 or (start is not None and start < 0):
return list(range(get_max_seg_number_cached(self) + 1))[s]
else:
return list(range(end + 1))[s]
def __str__(self) -> str:
return f"{self.dongle_id}/{self.log_id}" + (f"/{self.slice}" if self.slice else "") + (f"/{self.selector}" if self.selector else "")
def __repr__(self) -> str:
return self.__str__()

0
iqpilot/tools/lib/tests/__init__.py Normal file
View File

191
iqpilot/tools/lib/tests/test_caching.py Normal file
View File

@@ -0,0 +1,191 @@
import http.server
import multiprocessing
import os
import shutil
import socket
import tempfile
import pytest
from iqpilot.selfdrive.test.helpers import http_server_context
from iqpilot.system.hardware.hw import Paths
from iqpilot.tools.lib.url_file import URLFile, prune_cache
import iqpilot.tools.lib.url_file as url_file_module
def concurrent_prune_cache(cache_root, entry, barrier):
Paths.download_cache_root = staticmethod(lambda: cache_root)
barrier.wait()
prune_cache(entry)
class CachingTestRequestHandler(http.server.BaseHTTPRequestHandler):
FILE_EXISTS = True
def do_GET(self):
if self.FILE_EXISTS:
self.send_response(206 if "Range" in self.headers else 200, b'1234')
else:
self.send_response(404)
self.end_headers()
def do_HEAD(self):
if self.FILE_EXISTS:
self.send_response(200)
self.send_header("Content-Length", "4")
else:
self.send_response(404)
self.end_headers()
@pytest.fixture
def host():
with http_server_context(handler=CachingTestRequestHandler) as (host, port):
yield f"http://{host}:{port}"
class TestFileDownload:
def test_pipeline_defaults(self, host):
# TODO: parameterize the defaults so we don't rely on hard-coded values in xx
assert URLFile.pool_manager().pools._maxsize == 10# PoolManager num_pools param
pool_manager_defaults = {
"maxsize": 100,
"socket_options": [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1),],
}
for k, v in pool_manager_defaults.items():
assert URLFile.pool_manager().connection_pool_kw.get(k) == v
retry_defaults = {
"total": 6,
"backoff_factor": 0.75,
"status_forcelist": [409, 429, 500, 502, 503, 504],
}
for k, v in retry_defaults.items():
assert getattr(URLFile.pool_manager().connection_pool_kw["retries"], k) == v
# ensure caching on by default and cache dir gets created
os.environ.pop("DISABLE_FILEREADER_CACHE", None)
if os.path.exists(Paths.download_cache_root()):
shutil.rmtree(Paths.download_cache_root())
URLFile(f"{host}/test.txt").get_length()
URLFile(f"{host}/test.txt").read()
assert os.path.exists(Paths.download_cache_root())
def compare_loads(self, url, start=0, length=None):
"""Compares range between cached and non cached version"""
file_cached = URLFile(url, cache=True)
file_downloaded = URLFile(url, cache=False)
file_cached.seek(start)
file_downloaded.seek(start)
assert file_cached.get_length() == file_downloaded.get_length()
assert length + start if length is not None else 0 <= file_downloaded.get_length()
response_cached = file_cached.read(ll=length)
response_downloaded = file_downloaded.read(ll=length)
assert response_cached == response_downloaded
# Now test with cache in place
file_cached = URLFile(url, cache=True)
file_cached.seek(start)
response_cached = file_cached.read(ll=length)
assert file_cached.get_length() == file_downloaded.get_length()
assert response_cached == response_downloaded
def test_small_file(self):
# Make sure we don't force cache
os.environ.pop("DISABLE_FILEREADER_CACHE", None)
small_file_url = "https://raw.githubusercontent.com/commaai/openpilot/master/docs/SAFETY.md"
# If you want large file to be larger than a chunk
# large_file_url = "https://commadataci.blob.core.windows.net/openpilotci/0375fdf7b1ce594d/2019-06-13--08-32-25/3/fcamera.hevc"
# Load full small file
self.compare_loads(small_file_url)
file_small = URLFile(small_file_url)
length = file_small.get_length()
self.compare_loads(small_file_url, length - 100, 100)
self.compare_loads(small_file_url, 50, 100)
# Load small file 100 bytes at a time
for i in range(length // 100):
self.compare_loads(small_file_url, 100 * i, 100)
def test_large_file(self):
large_file_url = "https://commadataci.blob.core.windows.net/openpilotci/0375fdf7b1ce594d/2019-06-13--08-32-25/3/qlog.bz2"
# Load the end 100 bytes of both files
file_large = URLFile(large_file_url)
length = file_large.get_length()
self.compare_loads(large_file_url, length - 100, 100)
self.compare_loads(large_file_url)
@pytest.mark.parametrize("cache_enabled", [True, False])
def test_recover_from_missing_file(self, host, cache_enabled):
if cache_enabled:
os.environ.pop("DISABLE_FILEREADER_CACHE", None)
else:
os.environ["DISABLE_FILEREADER_CACHE"] = "1"
file_url = f"{host}/test.png"
CachingTestRequestHandler.FILE_EXISTS = False
length = URLFile(file_url).get_length()
assert length == -1
CachingTestRequestHandler.FILE_EXISTS = True
length = URLFile(file_url).get_length()
assert length == 4
class TestCache:
def test_concurrent_prune_cache(self, tmp_path):
context = multiprocessing.get_context("fork")
barrier = context.Barrier(16)
processes = [context.Process(target=concurrent_prune_cache, args=(f"{tmp_path}/", f"entry_{i}", barrier)) for i in range(16)]
for process in processes:
process.start()
for process in processes:
process.join(10)
assert process.exitcode == 0
manifest = set()
for line in (tmp_path / "manifest.txt").read_text().splitlines():
parts = line.split()
if len(parts) == 2:
manifest.add(parts[0])
assert manifest == {f"entry_{i}" for i in range(16)}
def test_prune_cache(self, monkeypatch):
with tempfile.TemporaryDirectory() as tmpdir:
monkeypatch.setattr(Paths, 'download_cache_root', staticmethod(lambda: tmpdir + "/"))
# setup test files and manifest
manifest_lines = []
for i in range(3):
fname = f"hash_{i}"
with open(tmpdir + "/" + fname, "wb") as f:
f.truncate(1000)
manifest_lines.append(f"{fname} {1000 + i}")
with open(tmpdir + "/manifest.txt", "w") as f:
f.write('\n'.join(manifest_lines))
# under limit, shouldn't prune
assert len(os.listdir(tmpdir)) == 4
prune_cache()
assert len([name for name in os.listdir(tmpdir) if name != "manifest.lock"]) == 4
# set a tiny cache limit to force eviction (1.5 chunks worth)
monkeypatch.setattr(url_file_module, 'CACHE_SIZE', url_file_module.CHUNK_SIZE + url_file_module.CHUNK_SIZE // 2)
# prune_cache should evict oldest files to get under limit
prune_cache()
remaining = [name for name in os.listdir(tmpdir) if name != "manifest.lock"]
# should have evicted at least one file + manifest
assert len(remaining) < 4
# newest file should remain
assert manifest_lines[2].split()[0] in remaining

176
iqpilot/tools/lib/tests/test_logreader.py Normal file
View File

@@ -0,0 +1,176 @@
import capnp
import contextlib
import shutil
import tempfile
import os
import pytest
import requests
from parameterized import parameterized
from iqpilot.cereal import log as capnp_log
from iqpilot.tools.lib.logreader import InternalUnavailableException, LogReader, parse_indirect
from iqpilot.tools.lib.route import SegmentRange
from iqpilot.tools.lib.url_file import URLFileException
NUM_SEGS = 17 # number of segments in the test route
ALL_SEGS = list(range(NUM_SEGS))
TEST_ROUTE = "344c5c15b34f2d8a/2024-01-03--09-37-12"
QLOG_FILE = "https://commadataci.blob.core.windows.net/openpilotci/0375fdf7b1ce594d/2019-06-13--08-32-25/3/qlog.bz2"
@contextlib.contextmanager
def setup_source_scenario(mocker, is_internal=False):
internal_source_mock = mocker.patch("iqpilot.tools.lib.logreader.internal_source")
internal_source_mock.__name__ = internal_source_mock._mock_name
openpilotci_source_mock = mocker.patch("iqpilot.tools.lib.logreader.openpilotci_source")
openpilotci_source_mock.__name__ = openpilotci_source_mock._mock_name
comma_api_source_mock = mocker.patch("iqpilot.tools.lib.logreader.comma_api_source")
comma_api_source_mock.__name__ = comma_api_source_mock._mock_name
if is_internal:
internal_source_mock.return_value = {3: QLOG_FILE}
else:
internal_source_mock.side_effect = InternalUnavailableException
openpilotci_source_mock.return_value = {}
comma_api_source_mock.return_value = {3: QLOG_FILE}
yield
class TestLogReader:
@pytest.mark.parametrize(("identifier", "expected"), [
(f"{TEST_ROUTE}", ALL_SEGS),
(f"{TEST_ROUTE.replace('/', '|')}", ALL_SEGS),
(f"{TEST_ROUTE}--0", [0]),
(f"{TEST_ROUTE}--5", [5]),
(f"{TEST_ROUTE}/0", [0]),
(f"{TEST_ROUTE}/5", [5]),
(f"{TEST_ROUTE}/0:10", ALL_SEGS[0:10]),
(f"{TEST_ROUTE}/0:0", []),
(f"{TEST_ROUTE}/4:6", ALL_SEGS[4:6]),
(f"{TEST_ROUTE}/0:-1", ALL_SEGS[0:-1]),
(f"{TEST_ROUTE}/:5", ALL_SEGS[:5]),
(f"{TEST_ROUTE}/2:", ALL_SEGS[2:]),
(f"{TEST_ROUTE}/2:-1", ALL_SEGS[2:-1]),
(f"{TEST_ROUTE}/-1", [ALL_SEGS[-1]]),
(f"{TEST_ROUTE}/-2", [ALL_SEGS[-2]]),
(f"{TEST_ROUTE}/-2:-1", ALL_SEGS[-2:-1]),
(f"{TEST_ROUTE}/-4:-2", ALL_SEGS[-4:-2]),
(f"{TEST_ROUTE}/:10:2", ALL_SEGS[:10:2]),
(f"{TEST_ROUTE}/5::2", ALL_SEGS[5::2]),
(f"https://useradmin.comma.ai/?onebox={TEST_ROUTE}", ALL_SEGS),
(f"https://useradmin.comma.ai/?onebox={TEST_ROUTE.replace('/', '|')}", ALL_SEGS),
(f"https://useradmin.comma.ai/?onebox={TEST_ROUTE.replace('/', '%7C')}", ALL_SEGS),
])
def test_indirect_parsing(self, identifier, expected, mocker):
mocker.patch("iqpilot.tools.lib.route.get_max_seg_number_cached", return_value=NUM_SEGS - 1)
parsed = parse_indirect(identifier)
sr = SegmentRange(parsed)
assert list(sr.seg_idxs) == expected, identifier
@parameterized.expand([
(f"{TEST_ROUTE}", f"{TEST_ROUTE}"),
(f"{TEST_ROUTE.replace('/', '|')}", f"{TEST_ROUTE}"),
(f"{TEST_ROUTE}--5", f"{TEST_ROUTE}/5"),
(f"{TEST_ROUTE}/0/q", f"{TEST_ROUTE}/0/q"),
(f"{TEST_ROUTE}/5:6/r", f"{TEST_ROUTE}/5:6/r"),
(f"{TEST_ROUTE}/5", f"{TEST_ROUTE}/5"),
])
def test_canonical_name(self, identifier, expected):
sr = SegmentRange(identifier)
assert str(sr) == expected
@pytest.mark.parametrize("cache_enabled", [True, False])
def test_direct_parsing(self, mocker, cache_enabled):
file_exists_mock = mocker.patch("iqpilot.tools.lib.filereader.file_exists")
if cache_enabled:
os.environ.pop("DISABLE_FILEREADER_CACHE", None)
else:
os.environ["DISABLE_FILEREADER_CACHE"] = "1"
qlog = tempfile.NamedTemporaryFile(mode='wb', delete=False)
with requests.get(QLOG_FILE, stream=True) as r:
with qlog as f:
shutil.copyfileobj(r.raw, f)
for f in [QLOG_FILE, qlog.name]:
l = len(list(LogReader(f)))
assert l > 100
with pytest.raises(URLFileException) if not cache_enabled else pytest.raises(AssertionError):
l = len(list(LogReader(QLOG_FILE.replace("/3/", "/200/"))))
# file_exists should not be called for direct files
assert file_exists_mock.call_count == 0
@parameterized.expand([
(f"{TEST_ROUTE}///",),
(f"{TEST_ROUTE}---",),
(f"{TEST_ROUTE}/-4:--2",),
(f"{TEST_ROUTE}/-a",),
(f"{TEST_ROUTE}/j",),
(f"{TEST_ROUTE}/0:1:2:3",),
(f"{TEST_ROUTE}/:::3",),
(f"{TEST_ROUTE}3",),
(f"{TEST_ROUTE}-3",),
(f"{TEST_ROUTE}--3a",),
])
def test_bad_ranges(self, segment_range):
with pytest.raises(AssertionError):
_ = SegmentRange(segment_range).seg_idxs
@pytest.mark.parametrize("segment_range, api_call", [
(f"{TEST_ROUTE}/0", False),
(f"{TEST_ROUTE}/:2", False),
(f"{TEST_ROUTE}/0:", True),
(f"{TEST_ROUTE}/-1", True),
(f"{TEST_ROUTE}", True),
])
def test_slicing_api_call(self, mocker, segment_range, api_call):
max_seg_mock = mocker.patch("iqpilot.tools.lib.route.get_max_seg_number_cached")
max_seg_mock.return_value = NUM_SEGS
_ = SegmentRange(segment_range).seg_idxs
assert api_call == max_seg_mock.called
@pytest.mark.parametrize("is_internal", [True, False])
def test_auto_source_scenarios(self, mocker, is_internal):
lr = LogReader(QLOG_FILE)
qlog_len = len(list(lr))
with setup_source_scenario(mocker, is_internal=is_internal):
lr = LogReader(f"{TEST_ROUTE}/3/q")
log_len = len(list(lr))
assert qlog_len == log_len
def test_only_union_types(self):
with tempfile.NamedTemporaryFile() as qlog:
# write valid Event messages
num_msgs = 100
with open(qlog.name, "wb") as f:
f.write(b"".join(capnp_log.Event.new_message().to_bytes() for _ in range(num_msgs)))
msgs = list(LogReader(qlog.name))
assert len(msgs) == num_msgs
[m.which() for m in msgs]
# append non-union Event message
event_msg = capnp_log.Event.new_message()
non_union_bytes = bytearray(event_msg.to_bytes())
non_union_bytes[event_msg.total_size.word_count * 8] = 0xff # set discriminant value out of range using Event word offset
with open(qlog.name, "ab") as f:
f.write(non_union_bytes)
# ensure new message is added, but is not a union type
msgs = list(LogReader(qlog.name))
assert len(msgs) == num_msgs + 1
with pytest.raises((capnp.KjException, RuntimeError)):
[m.which() for m in msgs]
# should not be added when only_union_types=True
msgs = list(LogReader(qlog.name, only_union_types=True))
assert len(msgs) == num_msgs
[m.which() for m in msgs]

27
iqpilot/tools/lib/tests/test_route_library.py Normal file
View File

@@ -0,0 +1,27 @@
from collections import namedtuple
from iqpilot.tools.lib.route import SegmentName
class TestRouteLibrary:
def test_segment_name_formats(self):
Case = namedtuple('Case', ['input', 'expected_route', 'expected_segment_num', 'expected_data_dir'])
cases = [ Case("a2a0ccea32023010|2023-07-27--13-01-19", "a2a0ccea32023010|2023-07-27--13-01-19", -1, None),
Case("a2a0ccea32023010/2023-07-27--13-01-19--1", "a2a0ccea32023010|2023-07-27--13-01-19", 1, None),
Case("a2a0ccea32023010|2023-07-27--13-01-19/2", "a2a0ccea32023010|2023-07-27--13-01-19", 2, None),
Case("a2a0ccea32023010/2023-07-27--13-01-19/3", "a2a0ccea32023010|2023-07-27--13-01-19", 3, None),
Case("/data/media/0/realdata/a2a0ccea32023010|2023-07-27--13-01-19", "a2a0ccea32023010|2023-07-27--13-01-19", -1, "/data/media/0/realdata"),
Case("/data/media/0/realdata/a2a0ccea32023010|2023-07-27--13-01-19--1", "a2a0ccea32023010|2023-07-27--13-01-19", 1, "/data/media/0/realdata"),
Case("/data/media/0/realdata/a2a0ccea32023010|2023-07-27--13-01-19/2", "a2a0ccea32023010|2023-07-27--13-01-19", 2, "/data/media/0/realdata") ]
def _validate(case):
route_or_segment_name = case.input
s = SegmentName(route_or_segment_name, allow_route_name=True)
assert str(s.route_name) == case.expected_route
assert s.segment_num == case.expected_segment_num
assert s.data_dir == case.expected_data_dir
for case in cases:
_validate(case)

244
iqpilot/tools/lib/url_file.py Normal file
View File

@@ -0,0 +1,244 @@
import logging
import fcntl
import os
import re
import socket
import time
from hashlib import md5
from urllib3 import PoolManager, Retry
from urllib3.response import BaseHTTPResponse
from urllib3.util import Timeout
from iqpilot.common.utils import atomic_write
from iqpilot.system.hardware.hw import Paths
from urllib3.exceptions import MaxRetryError
# Cache chunk size
K = 1000
CHUNK_SIZE = 1000 * K
CACHE_SIZE = 10 * 1024 * 1024 * 1024 # total cache size in GB
logging.getLogger("urllib3").setLevel(logging.WARNING)
def _env_int(name: str, default: int) -> int:
raw = os.getenv(name)
if raw is None:
return default
try:
value = int(raw)
return value if value > 0 else default
except ValueError:
return default
def hash_url(link: str) -> str:
return md5((link.split("?")[0]).encode('utf-8')).hexdigest()
def prune_cache(new_entry: str | None = None) -> None:
"""Evicts oldest cache files (LRU) until cache is under the size limit."""
cache_root = Paths.download_cache_root()
os.makedirs(cache_root, exist_ok=True)
manifest_path = os.path.join(cache_root, "manifest.txt")
with open(os.path.join(cache_root, "manifest.lock"), "w") as lock_file:
fcntl.flock(lock_file, fcntl.LOCK_EX)
manifest = {}
try:
with open(manifest_path) as f:
manifest = {parts[0]: int(parts[1]) for line in f if (parts := line.strip().split()) and len(parts) == 2}
except FileNotFoundError:
pass
if new_entry:
manifest[new_entry] = int(time.time()) # noqa: TID251
sorted_items = sorted(manifest.items(), key=lambda x: x[1])
while len(manifest) * CHUNK_SIZE > CACHE_SIZE and sorted_items:
key, _ = sorted_items.pop(0)
try:
os.remove(os.path.join(cache_root, key))
except OSError:
pass
manifest.pop(key, None)
with atomic_write(manifest_path, mode="w", overwrite=True) as f:
f.write('\n'.join(f"{k} {v}" for k, v in manifest.items()))
class URLFileException(Exception):
pass
class URLFile:
_pool_manager: PoolManager | None = None
@staticmethod
def reset() -> None:
URLFile._pool_manager = None
@staticmethod
def pool_manager() -> PoolManager:
if URLFile._pool_manager is None:
socket_options = [(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)]
retries = Retry(
total=_env_int("URLFILE_RETRIES_TOTAL", 6),
connect=_env_int("URLFILE_RETRIES_CONNECT", 6),
read=_env_int("URLFILE_RETRIES_READ", 6),
backoff_factor=float(os.getenv("URLFILE_RETRIES_BACKOFF", "0.75")),
status_forcelist=[409, 429, 500, 502, 503, 504],
)
URLFile._pool_manager = PoolManager(num_pools=10, maxsize=100, socket_options=socket_options, retries=retries)
return URLFile._pool_manager
def __init__(self, url: str, timeout: int = 10, cache: bool | None = None):
self._url = url
connect_timeout = _env_int("URLFILE_CONNECT_TIMEOUT", min(timeout, 10))
read_timeout = _env_int("URLFILE_READ_TIMEOUT", max(timeout, 30))
total_timeout = _env_int("URLFILE_TOTAL_TIMEOUT", max(read_timeout * 4, 180))
self._timeout = Timeout(connect=connect_timeout, read=read_timeout, total=total_timeout)
self._pos = 0
self._length: int | None = None
# Caching enabled by default, can be disabled with DISABLE_FILEREADER_CACHE=1, or overwritten by the cache input
self._force_download = int(os.environ.get("DISABLE_FILEREADER_CACHE", "0")) == 1
if cache is not None:
self._force_download = not cache
if not self._force_download:
os.makedirs(Paths.download_cache_root(), exist_ok=True)
def __enter__(self):
return self
def __exit__(self, exc_type, exc_value, traceback) -> None:
pass
def _request(self, method: str, url: str, headers: dict[str, str] | None = None) -> BaseHTTPResponse:
try:
return URLFile.pool_manager().request(method, url, timeout=self._timeout, headers=headers)
except MaxRetryError as e:
raise URLFileException(f"Failed to {method} {url}: {e}") from e
def get_length_online(self) -> int:
response = self._request('HEAD', self._url)
if not (200 <= response.status <= 299):
return -1
length = response.headers.get('content-length', 0)
return int(length)
def get_length(self) -> int:
if self._length is not None:
return self._length
file_length_path = os.path.join(Paths.download_cache_root(), hash_url(self._url) + "_length")
if not self._force_download and os.path.exists(file_length_path):
with open(file_length_path) as file_length:
content = file_length.read()
self._length = int(content)
return self._length
self._length = self.get_length_online()
if not self._force_download and self._length != -1:
with atomic_write(file_length_path, mode="w", overwrite=True) as file_length:
file_length.write(str(self._length))
return self._length
def read(self, ll: int | None = None) -> bytes:
if self._force_download:
return self.read_aux(ll=ll)
file_begin = self._pos
file_end = self._pos + ll if ll is not None else self.get_length()
assert file_end != -1, f"Remote file is empty or doesn't exist: {self._url}"
# We have to align with chunks we store. Position is the begginiing of the latest chunk that starts before or at our file
position = (file_begin // CHUNK_SIZE) * CHUNK_SIZE
response = b""
while True:
self._pos = position
chunk_number = self._pos / CHUNK_SIZE
file_name = hash_url(self._url) + "_" + str(chunk_number)
full_path = os.path.join(Paths.download_cache_root(), str(file_name))
data = None
# If we don't have a file, download it
if not os.path.exists(full_path):
data = self.read_aux(ll=CHUNK_SIZE)
with atomic_write(full_path, mode="wb", overwrite=True) as new_cached_file:
new_cached_file.write(data)
prune_cache(file_name)
else:
with open(full_path, "rb") as cached_file:
data = cached_file.read()
response += data[max(0, file_begin - position): min(CHUNK_SIZE, file_end - position)]
position += CHUNK_SIZE
if position >= file_end:
self._pos = file_end
return response
def read_aux(self, ll: int | None = None) -> bytes:
if ll is None:
length = self.get_length()
if length == -1:
raise URLFileException(f"Remote file is empty or doesn't exist: {self._url}")
end = length
else:
end = self._pos + ll
data = self.get_multi_range([(self._pos, end)])
self._pos += len(data[0])
return data[0]
def get_multi_range(self, ranges: list[tuple[int, int]]) -> list[bytes]:
# HTTP range requests are inclusive
assert all(e > s for s, e in ranges), "Range end must be greater than start"
rs = [f"{s}-{e-1}" for s, e in ranges if e > s]
r = self._request("GET", self._url, headers={"Range": "bytes=" + ",".join(rs)})
if r.status not in [200, 206]:
raise URLFileException(f"Expected 206 or 200 response {r.status} ({self._url})")
ctype = (r.headers.get("content-type") or "").lower()
if "multipart/byteranges" not in ctype:
return [r.data,]
m = re.search(r'boundary="?([^";]+)"?', ctype)
if not m:
raise URLFileException(f"Missing multipart boundary ({self._url})")
boundary = m.group(1).encode()
parts = []
for chunk in r.data.split(b"--" + boundary):
if b"\r\n\r\n" not in chunk:
continue
payload = chunk.split(b"\r\n\r\n", 1)[1].rstrip(b"\r\n")
if payload and payload != b"--":
parts.append(payload)
if len(parts) != len(ranges):
raise URLFileException(f"Expected {len(ranges)} parts, got {len(parts)} ({self._url})")
return parts
def seekable(self) -> bool:
return True
def seek(self, pos: int, whence: int = 0) -> int:
pos = int(pos)
if whence == os.SEEK_SET:
self._pos = pos
elif whence == os.SEEK_CUR:
self._pos += pos
elif whence == os.SEEK_END:
length = self.get_length()
assert length != -1, "Cannot seek from end on unknown length file"
self._pos = length + pos
else:
raise URLFileException("Invalid whence value")
return self._pos
def tell(self) -> int:
return self._pos
@property
def name(self) -> str:
return self._url
os.register_at_fork(after_in_child=URLFile.reset)

311
iqpilot/tools/lib/vidindex.py Executable file
View File

@@ -0,0 +1,311 @@
#!/usr/bin/env python3
import argparse
import os
import struct
from enum import IntEnum
from iqpilot.tools.lib.filereader import FileReader
DEBUG = int(os.getenv("DEBUG", "0"))
# compare to ffmpeg parsing
# ffmpeg -i <input.hevc> -c copy -bsf:v trace_headers -f null - 2>&1 | grep -B4 -A32 '] 0 '
# H.265 specification
# https://www.itu.int/rec/dologin_pub.asp?lang=e&id=T-REC-H.265-201802-S!!PDF-E&type=items
NAL_UNIT_START_CODE = b"\x00\x00\x01"
NAL_UNIT_START_CODE_SIZE = len(NAL_UNIT_START_CODE)
NAL_UNIT_HEADER_SIZE = 2
class HevcNalUnitType(IntEnum):
TRAIL_N = 0 # RBSP structure: slice_segment_layer_rbsp( )
TRAIL_R = 1 # RBSP structure: slice_segment_layer_rbsp( )
TSA_N = 2 # RBSP structure: slice_segment_layer_rbsp( )
TSA_R = 3 # RBSP structure: slice_segment_layer_rbsp( )
STSA_N = 4 # RBSP structure: slice_segment_layer_rbsp( )
STSA_R = 5 # RBSP structure: slice_segment_layer_rbsp( )
RADL_N = 6 # RBSP structure: slice_segment_layer_rbsp( )
RADL_R = 7 # RBSP structure: slice_segment_layer_rbsp( )
RASL_N = 8 # RBSP structure: slice_segment_layer_rbsp( )
RASL_R = 9 # RBSP structure: slice_segment_layer_rbsp( )
RSV_VCL_N10 = 10
RSV_VCL_R11 = 11
RSV_VCL_N12 = 12
RSV_VCL_R13 = 13
RSV_VCL_N14 = 14
RSV_VCL_R15 = 15
BLA_W_LP = 16 # RBSP structure: slice_segment_layer_rbsp( )
BLA_W_RADL = 17 # RBSP structure: slice_segment_layer_rbsp( )
BLA_N_LP = 18 # RBSP structure: slice_segment_layer_rbsp( )
IDR_W_RADL = 19 # RBSP structure: slice_segment_layer_rbsp( )
IDR_N_LP = 20 # RBSP structure: slice_segment_layer_rbsp( )
CRA_NUT = 21 # RBSP structure: slice_segment_layer_rbsp( )
RSV_IRAP_VCL22 = 22
RSV_IRAP_VCL23 = 23
RSV_VCL24 = 24
RSV_VCL25 = 25
RSV_VCL26 = 26
RSV_VCL27 = 27
RSV_VCL28 = 28
RSV_VCL29 = 29
RSV_VCL30 = 30
RSV_VCL31 = 31
VPS_NUT = 32 # RBSP structure: video_parameter_set_rbsp( )
SPS_NUT = 33 # RBSP structure: seq_parameter_set_rbsp( )
PPS_NUT = 34 # RBSP structure: pic_parameter_set_rbsp( )
AUD_NUT = 35
EOS_NUT = 36
EOB_NUT = 37
FD_NUT = 38
PREFIX_SEI_NUT = 39
SUFFIX_SEI_NUT = 40
RSV_NVCL41 = 41
RSV_NVCL42 = 42
RSV_NVCL43 = 43
RSV_NVCL44 = 44
RSV_NVCL45 = 45
RSV_NVCL46 = 46
RSV_NVCL47 = 47
UNSPEC48 = 48
UNSPEC49 = 49
UNSPEC50 = 50
UNSPEC51 = 51
UNSPEC52 = 52
UNSPEC53 = 53
UNSPEC54 = 54
UNSPEC55 = 55
UNSPEC56 = 56
UNSPEC57 = 57
UNSPEC58 = 58
UNSPEC59 = 59
UNSPEC60 = 60
UNSPEC61 = 61
UNSPEC62 = 62
UNSPEC63 = 63
# B.2.2 Byte stream NAL unit semantics
# - The nal_unit_type within the nal_unit( ) syntax structure is equal to VPS_NUT, SPS_NUT or PPS_NUT.
# - The byte stream NAL unit syntax structure contains the first NAL unit of an access unit in decoding
# order, as specified in clause 7.4.2.4.4.
HEVC_PARAMETER_SET_NAL_UNITS = (
HevcNalUnitType.VPS_NUT,
HevcNalUnitType.SPS_NUT,
HevcNalUnitType.PPS_NUT,
)
# 3.29 coded slice segment NAL unit: A NAL unit that has nal_unit_type in the range of TRAIL_N to RASL_R,
# inclusive, or in the range of BLA_W_LP to RSV_IRAP_VCL23, inclusive, which indicates that the NAL unit
# contains a coded slice segment
HEVC_CODED_SLICE_SEGMENT_NAL_UNITS = (
HevcNalUnitType.TRAIL_N,
HevcNalUnitType.TRAIL_R,
HevcNalUnitType.TSA_N,
HevcNalUnitType.TSA_R,
HevcNalUnitType.STSA_N,
HevcNalUnitType.STSA_R,
HevcNalUnitType.RADL_N,
HevcNalUnitType.RADL_R,
HevcNalUnitType.RASL_N,
HevcNalUnitType.RASL_R,
HevcNalUnitType.BLA_W_LP,
HevcNalUnitType.BLA_W_RADL,
HevcNalUnitType.BLA_N_LP,
HevcNalUnitType.IDR_W_RADL,
HevcNalUnitType.IDR_N_LP,
HevcNalUnitType.CRA_NUT,
)
class VideoFileInvalid(Exception):
pass
def get_ue(dat: bytes, start_idx: int, skip_bits: int) -> tuple[int, int]:
prefix_val = 0
prefix_len = 0
suffix_val = 0
suffix_len = 0
i = start_idx
while i < len(dat):
j = 7
while j >= 0:
if skip_bits > 0:
skip_bits -= 1
elif prefix_val == 0:
prefix_val = (dat[i] >> j) & 1
prefix_len += 1
else:
suffix_val = (suffix_val << 1) | ((dat[i] >> j) & 1)
suffix_len += 1
j -= 1
if prefix_val == 1 and prefix_len - 1 == suffix_len:
val = int(2**(prefix_len-1) - 1 + suffix_val)
size = prefix_len + suffix_len
return val, size
i += 1
raise VideoFileInvalid("invalid exponential-golomb code")
def require_nal_unit_start(dat: bytes, nal_unit_start: int) -> None:
if nal_unit_start < 1:
raise ValueError("start index must be greater than zero")
if dat[nal_unit_start:nal_unit_start + NAL_UNIT_START_CODE_SIZE] != NAL_UNIT_START_CODE:
raise VideoFileInvalid("data must begin with start code")
def get_hevc_nal_unit_length(dat: bytes, nal_unit_start: int) -> int:
try:
pos = dat.index(NAL_UNIT_START_CODE, nal_unit_start + NAL_UNIT_START_CODE_SIZE)
except ValueError:
pos = -1
# length of NAL unit is byte count up to next NAL unit start index
nal_unit_len = (pos if pos != -1 else len(dat)) - nal_unit_start
if DEBUG:
print(" nal_unit_len:", nal_unit_len)
return nal_unit_len
def get_hevc_nal_unit_type(dat: bytes, nal_unit_start: int) -> HevcNalUnitType:
# 7.3.1.2 NAL unit header syntax
# nal_unit_header( ) { // descriptor
# forbidden_zero_bit f(1)
# nal_unit_type u(6)
# nuh_layer_id u(6)
# nuh_temporal_id_plus1 u(3)
# }
header_start = nal_unit_start + NAL_UNIT_START_CODE_SIZE
nal_unit_header = dat[header_start:header_start + NAL_UNIT_HEADER_SIZE]
if len(nal_unit_header) != 2:
raise VideoFileInvalid("data to short to contain nal unit header")
nal_unit_type = HevcNalUnitType((nal_unit_header[0] >> 1) & 0x3F)
if DEBUG:
print(" nal_unit_type:", nal_unit_type.name, f"({nal_unit_type.value})")
return nal_unit_type
def get_hevc_slice_type(dat: bytes, nal_unit_start: int, nal_unit_type: HevcNalUnitType) -> tuple[int, bool]:
# 7.3.2.9 Slice segment layer RBSP syntax
# slice_segment_layer_rbsp( ) {
# slice_segment_header( )
# slice_segment_data( )
# rbsp_slice_segment_trailing_bits( )
# }
# ...
# 7.3.6.1 General slice segment header syntax
# slice_segment_header( ) { // descriptor
# first_slice_segment_in_pic_flag u(1)
# if( nal_unit_type >= BLA_W_LP && nal_unit_type <= RSV_IRAP_VCL23 )
# no_output_of_prior_pics_flag u(1)
# slice_pic_parameter_set_id ue(v)
# if( !first_slice_segment_in_pic_flag ) {
# if( dependent_slice_segments_enabled_flag )
# dependent_slice_segment_flag u(1)
# slice_segment_address u(v)
# }
# if( !dependent_slice_segment_flag ) {
# for( i = 0; i < num_extra_slice_header_bits; i++ )
# slice_reserved_flag[ i ] u(1)
# slice_type ue(v)
# ...
rbsp_start = nal_unit_start + NAL_UNIT_START_CODE_SIZE + NAL_UNIT_HEADER_SIZE
skip_bits = 0
# 7.4.7.1 General slice segment header semantics
# first_slice_segment_in_pic_flag equal to 1 specifies that the slice segment is the first slice segment of the picture in
# decoding order. first_slice_segment_in_pic_flag equal to 0 specifies that the slice segment is not the first slice segment
# of the picture in decoding order.
is_first_slice = dat[rbsp_start] >> 7 & 1 == 1
if not is_first_slice:
# TODO: parse dependent_slice_segment_flag and slice_segment_address and get real slice_type
# for now since we don't use it return -1 for slice_type
return (-1, is_first_slice)
skip_bits += 1 # skip past first_slice_segment_in_pic_flag
if nal_unit_type >= HevcNalUnitType.BLA_W_LP and nal_unit_type <= HevcNalUnitType.RSV_IRAP_VCL23:
# 7.4.7.1 General slice segment header semantics
# no_output_of_prior_pics_flag affects the output of previously-decoded pictures in the decoded picture buffer after the
# decoding of an IDR or a BLA picture that is not the first picture in the bitstream as specified in Annex C.
skip_bits += 1 # skip past no_output_of_prior_pics_flag
# 7.4.7.1 General slice segment header semantics
# slice_pic_parameter_set_id specifies the value of pps_pic_parameter_set_id for the PPS in use.
# The value of slice_pic_parameter_set_id shall be in the range of 0 to 63, inclusive.
_, size = get_ue(dat, rbsp_start, skip_bits)
skip_bits += size # skip past slice_pic_parameter_set_id
# 7.4.3.3.1 General picture parameter set RBSP semanal_unit_lenntics
# num_extra_slice_header_bits specifies the number of extra slice header bits that are present in the slice header RBSP
# for coded pictures referring to the PPS. The value of num_extra_slice_header_bits shall be in the range of 0 to 2, inclusive,
# in bitstreams conforming to this version of this Specification. Other values for num_extra_slice_header_bits are reserved
# for future use by ITU-T | ISO/IEC. However, decoders shall allow num_extra_slice_header_bits to have any value.
# TODO: get from PPS_NUT pic_parameter_set_rbsp( ) for corresponding slice_pic_parameter_set_id
num_extra_slice_header_bits = 0
skip_bits += num_extra_slice_header_bits
# 7.4.7.1 General slice segment header semantics
# slice_type specifies the coding type of the slice according to Table 7-7.
# Table 7-7 - Name association to slice_type
# slice_type | Name of slice_type
# 0 | B (B slice)
# 1 | P (P slice)
# 2 | I (I slice)
# unsigned integer 0-th order Exp-Golomb-coded syntax element with the left bit first
slice_type, _ = get_ue(dat, rbsp_start, skip_bits)
if DEBUG:
print(" slice_type:", slice_type, f"(first slice: {is_first_slice})")
if slice_type > 2:
raise VideoFileInvalid("slice_type must be 0, 1, or 2")
return slice_type, is_first_slice
def hevc_index(hevc_file_name: str, allow_corrupt: bool=False) -> tuple[list, int, bytes]:
with FileReader(hevc_file_name) as f:
dat = f.read()
if len(dat) < NAL_UNIT_START_CODE_SIZE + 1:
raise VideoFileInvalid("data is too short")
if dat[0] != 0x00:
raise VideoFileInvalid("first byte must be 0x00")
prefix_dat = b""
frame_types = list()
i = 1 # skip past first byte 0x00
try:
while i < len(dat):
require_nal_unit_start(dat, i)
nal_unit_len = get_hevc_nal_unit_length(dat, i)
nal_unit_type = get_hevc_nal_unit_type(dat, i)
if nal_unit_type in HEVC_PARAMETER_SET_NAL_UNITS:
prefix_dat += dat[i:i+nal_unit_len]
elif nal_unit_type in HEVC_CODED_SLICE_SEGMENT_NAL_UNITS:
slice_type, is_first_slice = get_hevc_slice_type(dat, i, nal_unit_type)
if is_first_slice:
frame_types.append((slice_type, i))
i += nal_unit_len
except Exception as e:
if not allow_corrupt:
raise
print(f"ERROR: NAL unit skipped @ {i}\n", str(e))
return frame_types, len(dat), prefix_dat
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("input_file", type=str)
parser.add_argument("output_prefix_file", type=str)
parser.add_argument("output_index_file", type=str)
args = parser.parse_args()
frame_types, dat_len, prefix_dat = hevc_index(args.input_file)
with open(args.output_prefix_file, "wb") as f:
f.write(prefix_dat)
with open(args.output_index_file, "wb") as f:
for ft, fp in frame_types:
f.write(struct.pack("<II", ft, fp))
f.write(struct.pack("<II", 0xFFFFFFFF, dat_len))
if __name__ == "__main__":
main()