pixeltable 0.4.16__py3-none-any.whl → 0.4.18__py3-none-any.whl

pixeltable/iterators/video.py

@@ -1,10 +1,10 @@
+import glob
 import logging
 import math
-import shutil
 import subprocess
 from fractions import Fraction
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any, Iterator, Literal, Optional
 
 import av
 import pandas as pd
@@ -14,6 +14,7 @@ import pixeltable as pxt
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 import pixeltable.utils.av as av_utils
+from pixeltable.env import Env
 from pixeltable.utils.local_store import TempStore
 
 from .base import ComponentIterator
@@ -237,75 +238,120 @@ class VideoSplitter(ComponentIterator):
     seconds.
 
     Args:
-        segment_duration: Video segment duration in seconds
-        overlap: Overlap between consecutive segments in seconds.
-        min_segment_duration: Drop the last segment if it is smaller than min_segment_duration
+        duration: Video segment duration in seconds
+        overlap: Overlap between consecutive segments in seconds. Only available for `mode='fast'`.
+        min_segment_duration: Drop the last segment if it is smaller than min_segment_duration.
+        mode: Segmentation mode:
+            - `'fast'`: Quick segmentation using stream copy (splits only at keyframes, approximate durations)
+            - `'accurate'`: Precise segmentation with re-encoding (exact durations, slower)
+        video_encoder: Video encoder to use. If not specified, uses the default encoder for the current platform.
+            Only available for `mode='accurate'`.
+        video_encoder_args: Additional arguments to pass to the video encoder. Only available for `mode='accurate'`.
     """
 
     # Input parameters
     video_path: Path
-    segment_duration: float
+    segment_duration: float | None
+    segment_times: list[float] | None
     overlap: float
     min_segment_duration: float
+    video_encoder: str | None
+    video_encoder_args: dict[str, Any] | None
 
     # Video metadata
     video_duration: float
     video_time_base: Fraction
     video_start_time: int
 
-    # position tracking
-    next_segment_start: float
-    next_segment_start_pts: int
+    output_iter: Iterator[dict[str, Any]]
 
-    def __init__(self, video: str, segment_duration: float, *, overlap: float = 0.0, min_segment_duration: float = 0.0):
-        assert segment_duration > 0.0
-        assert segment_duration >= min_segment_duration
-        assert overlap < segment_duration
+    def __init__(
+        self,
+        video: str,
+        *,
+        duration: float | None = None,
+        overlap: float | None = None,
+        min_segment_duration: float | None = None,
+        segment_times: list[float] | None = None,
+        mode: Literal['fast', 'accurate'] = 'accurate',
+        video_encoder: str | None = None,
+        video_encoder_args: dict[str, Any] | None = None,
+    ):
+        Env.get().require_binary('ffmpeg')
+        assert (duration is not None) != (segment_times is not None)
+        if segment_times is not None:
+            assert len(segment_times) > 0
+        if duration is not None:
+            assert duration > 0.0
+            assert duration >= min_segment_duration
+            assert overlap is None or overlap < duration
 
         video_path = Path(video)
         assert video_path.exists() and video_path.is_file()
 
-        if not shutil.which('ffmpeg'):
-            raise pxt.Error('ffmpeg is not installed or not in PATH. Please install ffmpeg to use VideoSplitter.')
-
         self.video_path = video_path
-        self.segment_duration = segment_duration
-        self.overlap = overlap
-        self.min_segment_duration = min_segment_duration
+        self.segment_duration = duration
+        self.overlap = overlap if overlap is not None else 0.0
+        self.min_segment_duration = min_segment_duration if min_segment_duration is not None else 0.0
+        self.segment_times = segment_times
+        self.video_encoder = video_encoder
+        self.video_encoder_args = video_encoder_args
 
         with av.open(str(video_path)) as container:
             video_stream = container.streams.video[0]
             self.video_time_base = video_stream.time_base
             self.video_start_time = video_stream.start_time or 0
 
-        self.next_segment_start = float(self.video_start_time * self.video_time_base)
-        self.next_segment_start_pts = self.video_start_time
+        self.output_iter = self.fast_iter() if mode == 'fast' else self.accurate_iter()
 
     @classmethod
     def input_schema(cls) -> dict[str, ts.ColumnType]:
         return {
             'video': ts.VideoType(nullable=False),
-            'segment_duration': ts.FloatType(nullable=False),
+            'duration': ts.FloatType(nullable=True),
             'overlap': ts.FloatType(nullable=True),
             'min_segment_duration': ts.FloatType(nullable=True),
+            'segment_times': ts.JsonType(nullable=True),
+            'mode': ts.StringType(nullable=False),
+            'video_encoder': ts.StringType(nullable=True),
+            'video_encoder_args': ts.JsonType(nullable=True),
         }
 
     @classmethod
     def output_schema(cls, *args: Any, **kwargs: Any) -> tuple[dict[str, ts.ColumnType], list[str]]:
-        param_names = ['segment_duration', 'overlap', 'min_segment_duration']
+        param_names = ['duration', 'overlap', 'min_segment_duration', 'segment_times']
         params = dict(zip(param_names, args))
         params.update(kwargs)
 
-        segment_duration = params['segment_duration']
-        min_segment_duration = params.get('min_segment_duration', 0.0)
-        overlap = params.get('overlap', 0.0)
-
-        if segment_duration <= 0.0:
-            raise excs.Error('segment_duration must be a positive number')
-        if segment_duration < min_segment_duration:
-            raise excs.Error('segment_duration must be at least min_segment_duration')
-        if overlap >= segment_duration:
-            raise excs.Error('overlap must be less than segment_duration')
+        segment_duration = params.get('duration')
+        segment_times = params.get('segment_times')
+        overlap = params.get('overlap')
+        min_segment_duration = params.get('min_segment_duration')
+        mode = params.get('mode', 'fast')
+
+        if segment_duration is None and segment_times is None:
+            raise excs.Error('Must specify either duration or segment_times')
+        if segment_duration is not None and segment_times is not None:
+            raise excs.Error('duration and segment_times cannot both be specified')
+        if segment_times is not None:
+            if len(segment_times) == 0:
+                raise excs.Error('segment_times cannot be empty')
+            if overlap is not None:
+                raise excs.Error('overlap cannot be specified with segment_times')
+        if segment_duration is not None:
+            if segment_duration <= 0.0:
+                raise excs.Error('duration must be a positive number')
+            if min_segment_duration is not None and segment_duration < min_segment_duration:
+                raise excs.Error('duration must be at least min_segment_duration')
+            if overlap is not None and overlap >= segment_duration:
+                raise excs.Error('overlap must be less than duration')
+        if mode == 'accurate' and overlap is not None:
+            raise excs.Error("Cannot specify overlap for mode='accurate'")
+        if mode == 'fast':
+            if params.get('video_encoder') is not None:
+                raise excs.Error("Cannot specify video_encoder for mode='fast'")
+            if params.get('video_encoder_args') is not None:
+                raise excs.Error("Cannot specify video_encoder_args for mode='fast'")
 
         return {
             'segment_start': ts.FloatType(nullable=False),
@@ -315,48 +361,109 @@ class VideoSplitter(ComponentIterator):
             'video_segment': ts.VideoType(nullable=False),
         }, []
 
-    def __next__(self) -> dict[str, Any]:
-        segment_path = str(TempStore.create_path(extension='.mp4'))
+    def fast_iter(self) -> Iterator[dict[str, Any]]:
+        segment_path: str = ''
         try:
-            cmd = av_utils.ffmpeg_clip_cmd(
-                str(self.video_path), segment_path, self.next_segment_start, self.segment_duration
-            )
-            _ = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            start_time = 0.0
+            start_pts = 0
+            segment_idx = 0
+            while True:
+                target_duration: float | None
+                if self.segment_duration is not None:
+                    target_duration = self.segment_duration
+                elif self.segment_times is not None and segment_idx < len(self.segment_times):
+                    target_duration = self.segment_times[segment_idx] - start_time
+                else:
+                    target_duration = None  # the rest of the video
+
+                segment_path = str(TempStore.create_path(extension='.mp4'))
+                cmd = av_utils.ffmpeg_clip_cmd(str(self.video_path), segment_path, start_time, target_duration)
+                _ = subprocess.run(cmd, capture_output=True, text=True, check=True)
+
+                # use the actual duration
+                segment_duration = av_utils.get_video_duration(segment_path)
+                if segment_duration - self.overlap == 0.0 or segment_duration < self.min_segment_duration:
+                    # we're done
+                    Path(segment_path).unlink()
+                    return
+
+                segment_end = start_time + segment_duration
+                segment_end_pts = start_pts + round(segment_duration / self.video_time_base)
+                result = {
+                    'segment_start': start_time,
+                    'segment_start_pts': start_pts,
+                    'segment_end': segment_end,
+                    'segment_end_pts': segment_end_pts,
+                    'video_segment': segment_path,
+                }
+                yield result
 
-            # use the actual duration
-            segment_duration = av_utils.get_video_duration(segment_path)
-            if segment_duration - self.overlap == 0.0:
-                # we're done
-                Path(segment_path).unlink()
-                raise StopIteration
+                start_time = segment_end - self.overlap
+                start_pts = segment_end_pts - round(self.overlap / self.video_time_base)
 
-            if segment_duration < self.min_segment_duration:
-                Path(segment_path).unlink()
-                raise StopIteration
+                segment_idx += 1
+                if self.segment_times is not None and segment_idx > len(self.segment_times):
+                    # We've created all segments including the final segment after the last segment_time
+                    break
 
-            segment_end = self.next_segment_start + segment_duration
-            segment_end_pts = self.next_segment_start_pts + round(segment_duration / self.video_time_base)
+        except subprocess.CalledProcessError as e:
+            if segment_path and Path(segment_path).exists():
+                Path(segment_path).unlink()
+            error_msg = f'ffmpeg failed with return code {e.returncode}'
+            if e.stderr:
+                error_msg += f': {e.stderr.strip()}'
+            raise pxt.Error(error_msg) from e
 
-            result = {
-                'segment_start': self.next_segment_start,
-                'segment_start_pts': self.next_segment_start_pts,
-                'segment_end': segment_end,
-                'segment_end_pts': segment_end_pts,
-                'video_segment': segment_path,
-            }
-            self.next_segment_start = segment_end - self.overlap
-            self.next_segment_start_pts = segment_end_pts - round(self.overlap / self.video_time_base)
+    def accurate_iter(self) -> Iterator[dict[str, Any]]:
+        base_path = TempStore.create_path(extension='')
+        # Use ffmpeg -f segment for accurate segmentation with re-encoding
+        output_pattern = f'{base_path}_segment_%04d.mp4'
+        cmd = av_utils.ffmpeg_segment_cmd(
+            str(self.video_path),
+            output_pattern,
+            segment_duration=self.segment_duration,
+            segment_times=self.segment_times,
+            video_encoder=self.video_encoder,
+            video_encoder_args=self.video_encoder_args,
+        )
 
-            return result
+        try:
+            _ = subprocess.run(cmd, capture_output=True, text=True, check=True)
+            output_paths = sorted(glob.glob(f'{base_path}_segment_*.mp4'))
+            # TODO: is this actually an error?
+            # if len(output_paths) == 0:
+            #     stderr_output = result.stderr.strip() if result.stderr is not None else ''
+            #     raise pxt.Error(
+            #         f'ffmpeg failed to create output files for commandline: {" ".join(cmd)}\n{stderr_output}'
+            #     )
+            start_time = 0.0
+            start_pts = 0
+            for segment_path in output_paths:
+                segment_duration = av_utils.get_video_duration(segment_path)
+                if segment_duration < self.min_segment_duration:
+                    Path(segment_path).unlink()
+                    return
+
+                result = {
+                    'segment_start': start_time,
+                    'segment_start_pts': start_pts,
+                    'segment_end': start_time + segment_duration,
+                    'segment_end_pts': start_pts + round(segment_duration / self.video_time_base),
+                    'video_segment': segment_path,
+                }
+                yield result
+                start_time += segment_duration
+                start_pts += round(segment_duration / self.video_time_base)
 
         except subprocess.CalledProcessError as e:
-            if Path(segment_path).exists():
-                Path(segment_path).unlink()
             error_msg = f'ffmpeg failed with return code {e.returncode}'
             if e.stderr:
                 error_msg += f': {e.stderr.strip()}'
             raise pxt.Error(error_msg) from e
 
+    def __next__(self) -> dict[str, Any]:
+        return next(self.output_iter)
+
     def close(self) -> None:
         pass
 

pixeltable/plan.py

@@ -93,18 +93,13 @@ class SampleClause:
     seed: Optional[int]
     stratify_exprs: Optional[list[exprs.Expr]]
 
-    # This seed value is used if one is not supplied
-    DEFAULT_SEED = 0
-
     # The version of the hashing algorithm used for ordering and fractional sampling.
     CURRENT_VERSION = 1
 
     def __post_init__(self) -> None:
-        """If no version was provided, provide the default version"""
+        # If no version was provided, provide the default version
         if self.version is None:
             self.version = self.CURRENT_VERSION
-        if self.seed is None:
-            self.seed = self.DEFAULT_SEED
 
     @property
     def is_stratified(self) -> bool:
@@ -1006,6 +1001,7 @@ class Planner:
             analyzer.window_fn_calls
         )
         ctx = exec.ExecContext(row_builder)
+
         combined_ordering = cls._create_combined_ordering(analyzer, verify_agg=is_python_agg)
         cls._verify_join_clauses(analyzer)