Merge pull request #1682 from papr/pts-export

Export pts in *_timstamps.csv file

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()