[common] Add operators for Timecode

Author: Breakthrough

scenedetect/_cli/commands.py

@@ -158,9 +158,11 @@ def list_scenes(
                     " |  %5d  | %11d | %s | %11d | %s |"
                     % (
                         i + 1,
-                        start_time.frame_num + 1,
+                        start_time.timecode.pts
+                        if start_time.timecode
+                        else start_time.frame_num + 1,
                         start_time.get_timecode(),
-                        end_time.frame_num,
+                        end_time.timecode.pts if end_time.timecode else end_time.frame_num,
                         end_time.get_timecode(),
                     )
                     for i, (start_time, end_time) in enumerate(scenes)

scenedetect/common.py

@@ -144,6 +144,96 @@ class Timecode:
     def seconds(self) -> float:
         return float(self.time_base * self.pts)
 
+    def _get_other_as_seconds(self, other: ty.Any) -> float:
+        if isinstance(other, Timecode):
+            return other.seconds
+        if isinstance(other, float):
+            return float(other)
+        raise TypeError(f"Unsupported type for comparison with Timecode: {type(other)}")
+
+    def __add__(self, other: ty.Union[float, "Timecode"]) -> "Timecode":
+        if isinstance(other, Timecode):
+            if self.time_base != other.time_base:
+                raise ValueError("Timecode instances require equal time_base for arithmetic.")
+            return Timecode(self.pts + other.pts, self.time_base)
+        if isinstance(other, float):
+            # Assume other is in seconds. Convert to pts.
+            pts_offset = round(other / self.time_base)
+            return Timecode(self.pts + pts_offset, self.time_base)
+        return NotImplemented
+
+    def __sub__(self, other: ty.Union[float, "Timecode"]) -> "Timecode":
+        if isinstance(other, Timecode):
+            if self.time_base != other.time_base:
+                raise ValueError("Timecode instances require equal time_base for arithmetic.")
+            return Timecode(self.pts - other.pts, self.time_base)
+        if isinstance(other, float):
+            # Assume other is in seconds. Convert to pts.
+            pts_offset = round(other / self.time_base)
+            return Timecode(self.pts - pts_offset, self.time_base)
+        return NotImplemented
+
+    def __eq__(self, other: ty.Any) -> bool:
+        if isinstance(other, Timecode):
+            if self.time_base == other.time_base:
+                return self.pts == other.pts
+        try:
+            return math.isclose(self.seconds, self._get_other_as_seconds(other))
+        except TypeError:
+            return NotImplemented
+
+    def __ne__(self, other: ty.Any) -> bool:
+        eq_result = self.__eq__(other)
+        return not eq_result if eq_result is not NotImplemented else NotImplemented
+
+    def __lt__(self, other: ty.Any) -> bool:
+        if isinstance(other, Timecode):
+            if self.time_base == other.time_base:
+                return self.pts < other.pts
+        try:
+            other_seconds = self._get_other_as_seconds(other)
+            if math.isclose(self.seconds, other_seconds):
+                return False
+            return self.seconds < other_seconds
+        except TypeError:
+            return NotImplemented
+
+    def __le__(self, other: ty.Any) -> bool:
+        if isinstance(other, Timecode):
+            if self.time_base == other.time_base:
+                return self.pts <= other.pts
+        try:
+            other_seconds = self._get_other_as_seconds(other)
+            if math.isclose(self.seconds, other_seconds):
+                return True
+            return self.seconds < other_seconds
+        except TypeError:
+            return NotImplemented
+
+    def __gt__(self, other: ty.Any) -> bool:
+        if isinstance(other, Timecode):
+            if self.time_base == other.time_base:
+                return self.pts > other.pts
+        try:
+            other_seconds = self._get_other_as_seconds(other)
+            if math.isclose(self.seconds, other_seconds):
+                return False
+            return self.seconds > other_seconds
+        except TypeError:
+            return NotImplemented
+
+    def __ge__(self, other: ty.Any) -> bool:
+        if isinstance(other, Timecode):
+            if self.time_base == other.time_base:
+                return self.pts >= other.pts
+        try:
+            other_seconds = self._get_other_as_seconds(other)
+            if math.isclose(self.seconds, other_seconds):
+                return True
+            return self.seconds > other_seconds
+        except TypeError:
+            return NotImplemented
+
 
 class FrameTimecode:
     """Object for frame-based timecodes, using the video framerate to compute back and
@@ -211,9 +301,12 @@ def __init__(
         else:
             self._frame_num = self._parse_timecode_number(timecode)
 
-    # TODO(v0.7): Add a PTS property as well and slowly transition over to that, since we don't
-    # always know the position as a "frame number".  However, for the reverse case, we CAN state
-    # the presentation time if we know the frame number (for a fixed framerate video).
+    @property
+    def timecode(self) -> ty.Optional[Timecode]:
+        return self._timecode
+
+    # TODO(0.7): Can we make this return a Timecode if this is a VFR video? This value is used in
+    # a *lot* of places, so a large increase will be signficiant for things like buffer sizes.
     @property
     def frame_num(self) -> ty.Optional[int]:
         return self._frame_num
@@ -427,47 +520,69 @@ def __eq__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> bool:
         if other is None:
             return False
         if self._timecode:
-            return self.seconds == self._get_other_as_seconds(other)
+            if isinstance(other, FrameTimecode) and other._timecode:
+                return self._timecode == other._timecode
+            return math.isclose(self.seconds, self._get_other_as_seconds(other))
         return self.frame_num == self._get_other_as_frames(other)
 
     def __ne__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> bool:
         if other is None:
             return True
-        if self._timecode:
-            return self.seconds != self._get_other_as_seconds(other)
-        return self.frame_num != self._get_other_as_frames(other)
+        return not self.__eq__(other)
 
     def __lt__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> bool:
         if self._timecode:
-            return self.seconds < self._get_other_as_seconds(other)
+            if isinstance(other, FrameTimecode) and other._timecode:
+                return self._timecode < other._timecode
+            other_seconds = self._get_other_as_seconds(other)
+            if math.isclose(self.seconds, other_seconds):
+                return False
+            return self.seconds < other_seconds
         return self.frame_num < self._get_other_as_frames(other)
 
     def __le__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> bool:
         if self._timecode:
-            return self.seconds <= self._get_other_as_seconds(other)
+            if isinstance(other, FrameTimecode) and other._timecode:
+                return self._timecode <= other._timecode
+            other_seconds = self._get_other_as_seconds(other)
+            if math.isclose(self.seconds, other_seconds):
+                return True
+            return self.seconds < other_seconds
         return self.frame_num <= self._get_other_as_frames(other)
 
     def __gt__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> bool:
         if self._timecode:
-            return self.seconds > self._get_other_as_seconds(other)
+            if isinstance(other, FrameTimecode) and other._timecode:
+                return self._timecode > other._timecode
+            other_seconds = self._get_other_as_seconds(other)
+            if math.isclose(self.seconds, other_seconds):
+                return False
+            return self.seconds > other_seconds
         return self.frame_num > self._get_other_as_frames(other)
 
     def __ge__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> bool:
         if self._timecode:
-            return self.seconds >= self._get_other_as_seconds(other)
+            if isinstance(other, FrameTimecode) and other._timecode:
+                return self._timecode >= other._timecode
+            other_seconds = self._get_other_as_seconds(other)
+            if math.isclose(self.seconds, other_seconds):
+                return True
+            return self.seconds > other_seconds
         return self.frame_num >= self._get_other_as_frames(other)
 
     def __iadd__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> "FrameTimecode":
         if self._timecode:
-            new_seconds = self.seconds + self._get_other_as_seconds(other)
-            # TODO: This is incorrect for VFR, need a better way to handle this.
-            # For now, we convert back to a frame number.
-            self._frame_num = self._seconds_to_frames(new_seconds)
-            self._timecode = None
+            if isinstance(other, FrameTimecode) and other._timecode:
+                self._timecode = self._timecode + other._timecode
+            else:
+                other_as_seconds = self._get_other_as_seconds(other)
+                self._timecode = self._timecode + other_as_seconds
+            if self._timecode.pts < 0:
+                self._timecode = Timecode(0, self._timecode.time_base)
         else:
             self._frame_num += self._get_other_as_frames(other)
-        if self._frame_num < 0:  # Required to allow adding negative seconds/frames.
-            self._frame_num = 0
+            if self._frame_num < 0:  # Required to allow adding negative seconds/frames.
+                self._frame_num = 0
         return self
 
     def __add__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> "FrameTimecode":
@@ -477,15 +592,17 @@ def __add__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> "FrameTi
 
     def __isub__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> "FrameTimecode":
         if self._timecode:
-            new_seconds = self.seconds - self._get_other_as_seconds(other)
-            # TODO: This is incorrect for VFR, need a better way to handle this.
-            # For now, we convert back to a frame number.
-            self._frame_num = self._seconds_to_frames(new_seconds)
-            self._timecode = None
+            if isinstance(other, FrameTimecode) and other._timecode:
+                self._timecode = self._timecode - other._timecode
+            else:
+                other_as_seconds = self._get_other_as_seconds(other)
+                self._timecode = self._timecode - other_as_seconds
+            if self._timecode.pts < 0:
+                self._timecode = Timecode(0, self._timecode.time_base)
         else:
             self._frame_num -= self._get_other_as_frames(other)
-        if self._frame_num < 0:
-            self._frame_num = 0
+            if self._frame_num < 0:
+                self._frame_num = 0
         return self
 
     def __sub__(self, other: ty.Union[int, float, str, "FrameTimecode"]) -> "FrameTimecode":

scenedetect/output/__init__.py

@@ -98,13 +98,13 @@ def write_scene_list(
         csv_writer.writerow(
             [
                 "%d" % (i + 1),
-                "%d" % (start.frame_num + 1),
+                "%d" % (start.timecode.pts if start.timecode else start.frame_num + 1),
                 start.get_timecode(),
                 "%.3f" % start.seconds,
-                "%d" % end.frame_num,
+                "%d" % (end.timecode.pts if end.timecode else end.frame_num + 1),
                 end.get_timecode(),
                 "%.3f" % end.seconds,
-                "%d" % duration.frame_num,
+                "%d" % (duration.timecode.pts if duration.timecode else duration.frame_num + 1),
                 duration.get_timecode(),
                 "%.3f" % duration.seconds,
             ]

scenedetect/scene_manager.py

@@ -398,7 +398,7 @@ def _process_frame(
             cuts = detector.process_frame(position, frame_im)
             self._cutting_list += cuts
             new_cuts = True if cuts else False
-            # TODO: Support callbacks with PTS.
+            # TODO(0.7): Support callbacks with PTS.
             if callback:
                 if _USE_PTS_IN_DEVELOPMENT:
                     raise NotImplementedError()