diff --git a/pupil_src/shared_modules/av_writer.py b/pupil_src/shared_modules/av_writer.py index 4d0eeae537..a62de2c08f 100644 --- a/pupil_src/shared_modules/av_writer.py +++ b/pupil_src/shared_modules/av_writer.py @@ -9,20 +9,21 @@ ---------------------------------------------------------------------------~(*) """ +import abc + # logging import logging import multiprocessing as mp import os -import platform -from fractions import Fraction -import abc import typing as T +from fractions import Fraction +import av import numpy as np +from av.packet import Packet import audio_utils -import av -from av.packet import Packet +from video_capture.utils import Video logger = logging.getLogger(__name__) @@ -73,7 +74,19 @@ def write_timestamps(file_loc, timestamps, output_format="npy"): if output_format in ("npy", "all"): np.save(ts_loc + ".npy", ts) if output_format in ("csv", "all"): - np.savetxt(ts_loc + ".csv", ts, fmt="%f", header="timestamps [seconds]") + output_video = Video(file_loc) + if not output_video.is_valid: + logger.error(f"Failed to extract PTS frome exported video {file_loc}") + return + pts = output_video.pts + ts_pts = np.vstack((ts, pts)).T + np.savetxt( + ts_loc + ".csv", + ts_pts, + fmt=["%f", "%i"], + delimiter=",", + header="timestamps [seconds],pts", + ) class AV_Writer(abc.ABC): @@ -161,16 +174,21 @@ def write_video_frame(self, input_frame): # ensure strong monotonic pts pts = max(pts, self.last_video_pts + 1) - n_frames = self.video_stream.frames # TODO: Use custom Frame wrapper class, that wraps backend-specific frames. # This way we could just attach the pts here to the frame. # Currently this will fail e.g. for av.VideoFrame. + video_packed_encoded = False for packet in self.encode_frame(input_frame, pts): + if packet.stream is self.video_stream: + if video_packed_encoded: + # NOTE: Assumption: Each frame is encoded into a single packet! + # This is required for the frame.pts == packet.pts assumption below. + logger.warning("Single frame yielded more than one packet") + video_packed_encoded = True self.container.mux(packet) - if self.video_stream.frames != n_frames + 1: - logger.warning( - f"Failed to write frame #{n_frames} to video file. Ignoring." - ) + + if not video_packed_encoded: + logger.warning(f"Encoding frame {input_frame.index} failed!") return self.last_video_pts = pts @@ -188,14 +206,16 @@ def close(self, timestamp_export_format="npy"): for packet in self.video_stream.encode(None): self.container.mux(packet) - if timestamp_export_format is not None: - write_timestamps( - self.output_file_path, self.timestamps, timestamp_export_format - ) - self.container.close() self.closed = True + if self.configured and timestamp_export_format is not None: + # Requires self.container to be closed since we extract pts + # from the exported video file. + write_timestamps( + self.output_file_path, self.timestamps, timestamp_export_format + ) + def release(self): """Close writer, triggering stream and timestamp save.""" self.close()