framereader: add cache_dir argument (#29904)
* framereader: add cache_dir argument * make it an env var old-commit-hash: 880157f5a614a4a524273a9833e0a066fdec371e
This commit is contained in:
@@ -2,10 +2,10 @@ import os
|
||||
import urllib.parse
|
||||
from openpilot.common.file_helpers import mkdirs_exists_ok
|
||||
|
||||
DEFAULT_CACHE_DIR = os.path.expanduser("~/.commacache")
|
||||
DEFAULT_CACHE_DIR = os.getenv("CACHE_ROOT", os.path.expanduser("~/.commacache"))
|
||||
|
||||
def cache_path_for_file_path(fn, cache_prefix=None):
|
||||
dir_ = os.path.join(DEFAULT_CACHE_DIR, "local")
|
||||
def cache_path_for_file_path(fn, cache_dir=DEFAULT_CACHE_DIR):
|
||||
dir_ = os.path.join(cache_dir, "local")
|
||||
mkdirs_exists_ok(dir_)
|
||||
fn_parsed = urllib.parse.urlparse(fn)
|
||||
if fn_parsed.scheme == '':
|
||||
|
||||
@@ -12,7 +12,7 @@ import numpy as np
|
||||
from lru import LRU
|
||||
|
||||
import _io
|
||||
from openpilot.tools.lib.cache import cache_path_for_file_path
|
||||
from openpilot.tools.lib.cache import cache_path_for_file_path, DEFAULT_CACHE_DIR
|
||||
from openpilot.tools.lib.exceptions import DataUnreadableError
|
||||
from openpilot.common.file_helpers import atomic_write_in_dir
|
||||
|
||||
@@ -106,8 +106,8 @@ def cache_fn(func):
|
||||
if kwargs.pop('no_cache', None):
|
||||
cache_path = None
|
||||
else:
|
||||
cache_prefix = kwargs.pop('cache_prefix', None)
|
||||
cache_path = cache_path_for_file_path(fn, cache_prefix)
|
||||
cache_dir = kwargs.pop('cache_dir', DEFAULT_CACHE_DIR)
|
||||
cache_path = cache_path_for_file_path(fn, cache_dir)
|
||||
|
||||
if cache_path and os.path.exists(cache_path):
|
||||
with open(cache_path, "rb") as cache_file:
|
||||
@@ -140,18 +140,18 @@ def index_stream(fn, typ):
|
||||
}
|
||||
|
||||
|
||||
def index_videos(camera_paths, cache_prefix=None):
|
||||
def index_videos(camera_paths, cache_dir=DEFAULT_CACHE_DIR):
|
||||
"""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])
|
||||
for fn in camera_paths:
|
||||
index_video(fn, frame_type, cache_prefix)
|
||||
index_video(fn, frame_type, cache_dir)
|
||||
|
||||
|
||||
def index_video(fn, frame_type=None, cache_prefix=None):
|
||||
cache_path = cache_path_for_file_path(fn, cache_prefix)
|
||||
def index_video(fn, frame_type=None, cache_dir=DEFAULT_CACHE_DIR):
|
||||
cache_path = cache_path_for_file_path(fn, cache_dir)
|
||||
|
||||
if os.path.exists(cache_path):
|
||||
return
|
||||
@@ -160,16 +160,16 @@ def index_video(fn, frame_type=None, cache_prefix=None):
|
||||
frame_type = fingerprint_video(fn[0])
|
||||
|
||||
if frame_type == FrameType.h265_stream:
|
||||
index_stream(fn, "hevc", cache_prefix=cache_prefix)
|
||||
index_stream(fn, "hevc", cache_dir=cache_dir)
|
||||
else:
|
||||
raise NotImplementedError("Only h265 supported")
|
||||
|
||||
|
||||
def get_video_index(fn, frame_type, cache_prefix=None):
|
||||
cache_path = cache_path_for_file_path(fn, cache_prefix)
|
||||
def get_video_index(fn, frame_type, cache_dir=DEFAULT_CACHE_DIR):
|
||||
cache_path = cache_path_for_file_path(fn, cache_dir)
|
||||
|
||||
if not os.path.exists(cache_path):
|
||||
index_video(fn, frame_type, cache_prefix)
|
||||
index_video(fn, frame_type, cache_dir)
|
||||
|
||||
if not os.path.exists(cache_path):
|
||||
return None
|
||||
@@ -284,13 +284,13 @@ class BaseFrameReader:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
def FrameReader(fn, cache_prefix=None, readahead=False, readbehind=False, index_data=None):
|
||||
def FrameReader(fn, cache_dir=DEFAULT_CACHE_DIR, readahead=False, readbehind=False, index_data=None):
|
||||
frame_type = fingerprint_video(fn)
|
||||
if frame_type == FrameType.raw:
|
||||
return RawFrameReader(fn)
|
||||
elif frame_type in (FrameType.h265_stream,):
|
||||
if not index_data:
|
||||
index_data = get_video_index(fn, frame_type, cache_prefix)
|
||||
index_data = get_video_index(fn, frame_type, cache_dir)
|
||||
return StreamFrameReader(fn, frame_type, index_data, readahead=readahead, readbehind=readbehind)
|
||||
else:
|
||||
raise NotImplementedError(frame_type)
|
||||
|
||||
Reference in New Issue
Block a user