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
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 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 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
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 ffprobe(fn, fmt=None):
cmd = [
"ffprobe", "-v", "quiet", "-print_format", "json", "-show_format",
"-show_streams"
]
cmd += [fn]
try:
ffprobe_output = subprocess.check_output(cmd)
except subprocess.CalledProcessError as e:
raise DataUnreadableError(fn)
return json.loads(ffprobe_output)
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
else:
raise NotImplementedError(fn)
else:
raise NotImplementedError(fn)
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 decompress_video_data(rawdat, vid_fmt, w, h, pix_fmt):
# using a tempfile is much faster than proc.communicate for some reason
with tempfile.TemporaryFile() as tmpf:
tmpf.write(rawdat)
tmpf.seek(0)
threads = os.getenv("FFMPEG_THREADS", "0")
cuda = os.getenv("FFMPEG_CUDA", "0") == "1"
proc = subprocess.Popen(
["ffmpeg",
"-threads", threads,
"-hwaccel", "none" if not cuda else "cuda",
"-c:v", "hevc",
"-vsync", "0",
"-f", vid_fmt,
"-flags2", "showall",
"-i", "pipe:0",
"-threads", threads,
"-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 == "nv12":
ret = np.frombuffer(dat, dtype=np.uint8).reshape(-1, (h*w*3//2))
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
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
}
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 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])
