pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/io/utils.py

@@ -0,0 +1,100 @@
+from keyword import iskeyword as is_python_keyword
+from typing import Any
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable.catalog.globals import is_system_column_name
+
+
+def normalize_pxt_col_name(name: str) -> str:
+    """
+    Normalizes an arbitrary column name into a valid Pixeltable identifier by:
+    - replacing any non-ascii or non-alphanumeric characters with an underscore _
+    - prefixing the result with the letter 'c' if it starts with an underscore or a number
+    """
+    id = ''.join(ch if ch.isascii() and ch.isalnum() else '_' for ch in name)
+    if id[0].isnumeric():
+        id = f'c_{id}'
+    elif id[0] == '_':
+        id = f'c{id}'
+    assert pxt.catalog.is_valid_identifier(id), id
+    return id
+
+
+def normalize_primary_key_parameter(primary_key: str | list[str] | None = None) -> list[str]:
+    if primary_key is None:
+        primary_key = []
+    elif isinstance(primary_key, str):
+        primary_key = [primary_key]
+    elif not isinstance(primary_key, list) or not all(isinstance(pk, str) for pk in primary_key):
+        raise excs.Error('primary_key must be a single column name or a list of column names')
+    return primary_key
+
+
+def _is_usable_as_column_name(name: str, destination_schema: dict[str, Any]) -> bool:
+    return not (is_system_column_name(name) or is_python_keyword(name) or name in destination_schema)
+
+
+def normalize_schema_names(
+    in_schema: dict[str, Any],
+    primary_key: list[str],
+    schema_overrides: dict[str, Any],
+    require_valid_pxt_column_names: bool = False,
+) -> tuple[dict[str, Any], list[str], dict[str, str] | None]:
+    """
+    Convert all names in the input schema from source names to valid Pixeltable identifiers
+    - Ensure that all names are unique.
+    - Report an error if any types are missing
+    - If "require_valid_pxt_column_names", report an error if any column names are not valid Pixeltable column names
+    - Report an error if any primary key columns are missing
+    Returns
+    - A new schema with normalized column names
+    - The primary key columns, mapped to the normalized names
+    - A mapping from the original names to the normalized names.
+    """
+
+    # Report any untyped columns as an error
+    untyped_cols = [in_name for in_name, column_type in in_schema.items() if column_type is None]
+    if len(untyped_cols) > 0:
+        raise excs.Error(f'Could not infer pixeltable type for column(s): {", ".join(untyped_cols)}')
+
+    # Report any columns in `schema_overrides` that are not in the source
+    extraneous_overrides = schema_overrides.keys() - in_schema.keys()
+    if len(extraneous_overrides) > 0:
+        raise excs.Error(
+            f'Some column(s) specified in `schema_overrides` are not present '
+            f'in the source: {", ".join(extraneous_overrides)}'
+        )
+
+    schema: dict[str, Any] = {}
+    col_mapping: dict[str, str] = {}  # Maps column names to Pixeltable column names if needed
+    for in_name, pxt_type in in_schema.items():
+        pxt_name = normalize_pxt_col_name(in_name)
+        # Ensure that column names are unique by appending a distinguishing suffix
+        # to any collisions
+        pxt_fname = pxt_name
+        n = 1
+        while not _is_usable_as_column_name(pxt_fname, schema):
+            pxt_fname = f'{pxt_name}_{n}'
+            n += 1
+        schema[pxt_fname] = pxt_type
+        col_mapping[in_name] = pxt_fname
+
+    # Determine if the col_mapping is the identity mapping
+    non_identity_keys = [k for k, v in col_mapping.items() if k != v]
+    if len(non_identity_keys) > 0:
+        if require_valid_pxt_column_names:
+            raise excs.Error(
+                f'Column names must be valid pixeltable identifiers. Invalid names: {", ".join(non_identity_keys)}'
+            )
+    else:
+        col_mapping = None
+
+    # Report any primary key columns that are not in the source as an error
+    missing_pk = [pk for pk in primary_key if pk not in in_schema]
+    if len(missing_pk) > 0:
+        raise excs.Error(f'Primary key column(s) are not found in the source: {", ".join(missing_pk)}')
+
+    pxt_pk = [col_mapping[pk] for pk in primary_key] if col_mapping is not None else primary_key
+
+    return schema, pxt_pk, col_mapping

pixeltable/iterators/__init__.py

@@ -1,13 +1,17 @@
+"""Iterators for splitting media and documents into components."""
+# ruff: noqa: F401
+
+from .audio import AudioSplitter
 from .base import ComponentIterator
 from .document import DocumentSplitter
 from .image import TileIterator
 from .string import StringSplitter
-from .video import FrameIterator
+from .video import FrameIterator, VideoSplitter
 
-__default_dir = set(symbol for symbol in dir() if not symbol.startswith('_'))
+__default_dir = {symbol for symbol in dir() if not symbol.startswith('_')}
 __removed_symbols = {'base', 'document', 'video'}
-__all__ = sorted(list(__default_dir - __removed_symbols))
+__all__ = sorted(__default_dir - __removed_symbols)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/iterators/audio.py

@@ -0,0 +1,207 @@
+import logging
+from fractions import Fraction
+from pathlib import Path
+from typing import Any, ClassVar
+
+import av
+from deprecated import deprecated
+
+from pixeltable import exceptions as excs, type_system as ts
+from pixeltable.utils.local_store import TempStore
+
+from .base import ComponentIterator
+
+_logger = logging.getLogger('pixeltable')
+
+
+class AudioSplitter(ComponentIterator):
+    # Input parameters
+    audio_path: Path
+    chunk_duration_sec: float
+    overlap_sec: float
+
+    # audio stream details
+    container: av.container.input.InputContainer
+    audio_time_base: Fraction  # seconds per presentation time
+
+    # List of chunks to extract
+    # Each chunk is defined by start and end presentation timestamps in audio file (int)
+    chunks_to_extract_in_pts: list[tuple[int, int]] | None
+    # next chunk to extract
+    next_pos: int
+
+    __codec_map: ClassVar[dict[str, str]] = {
+        'mp3': 'mp3',  # MP3 decoder -> mp3/libmp3lame encoder
+        'mp3float': 'mp3',  # MP3float decoder -> mp3 encoder
+        'aac': 'aac',  # AAC decoder -> AAC encoder
+        'vorbis': 'libvorbis',  # Vorbis decoder -> libvorbis encoder
+        'opus': 'libopus',  # Opus decoder -> libopus encoder
+        'flac': 'flac',  # FLAC decoder -> FLAC encoder
+        'wavpack': 'wavpack',  # WavPack decoder -> WavPack encoder
+        'alac': 'alac',  # ALAC decoder -> ALAC encoder
+    }
+
+    def __init__(
+        self, audio: str, chunk_duration_sec: float, *, overlap_sec: float = 0.0, min_chunk_duration_sec: float = 0.0
+    ):
+        assert chunk_duration_sec > 0.0
+        assert chunk_duration_sec >= min_chunk_duration_sec
+        assert overlap_sec < chunk_duration_sec
+        audio_path = Path(audio)
+        assert audio_path.exists() and audio_path.is_file()
+        self.audio_path = audio_path
+        self.next_pos = 0
+        self.container = av.open(str(audio_path))
+        if len(self.container.streams.audio) == 0:
+            # No audio stream
+            return
+        self.chunk_duration_sec = chunk_duration_sec
+        self.overlap_sec = overlap_sec
+        self.min_chunk_duration_sec = min_chunk_duration_sec
+        self.audio_time_base = self.container.streams.audio[0].time_base
+
+        audio_start_time_pts = self.container.streams.audio[0].start_time or 0
+        audio_start_time_sec = float(audio_start_time_pts * self.audio_time_base)
+        total_audio_duration_pts = self.container.streams.audio[0].duration or 0
+        total_audio_duration_sec = float(total_audio_duration_pts * self.audio_time_base)
+
+        self.chunks_to_extract_in_pts = [
+            (round(start / self.audio_time_base), round(end / self.audio_time_base))
+            for (start, end) in self.build_chunks(
+                audio_start_time_sec, total_audio_duration_sec, chunk_duration_sec, overlap_sec, min_chunk_duration_sec
+            )
+        ]
+        _logger.debug(
+            f'AudioIterator: path={self.audio_path} total_audio_duration_pts={total_audio_duration_pts} '
+            f'chunks_to_extract_in_pts={self.chunks_to_extract_in_pts}'
+        )
+
+    @classmethod
+    def build_chunks(
+        cls,
+        start_time_sec: float,
+        total_duration_sec: float,
+        chunk_duration_sec: float,
+        overlap_sec: float,
+        min_chunk_duration_sec: float,
+    ) -> list[tuple[float, float]]:
+        chunks_to_extract_in_sec: list[tuple[float, float]] = []
+        current_pos = start_time_sec
+        end_time = start_time_sec + total_duration_sec
+        while current_pos < end_time:
+            chunk_start = current_pos
+            chunk_end = min(chunk_start + chunk_duration_sec, end_time)
+            chunks_to_extract_in_sec.append((chunk_start, chunk_end))
+            if chunk_end >= end_time:
+                break
+            current_pos = chunk_end - overlap_sec
+        # If the last chunk is smaller than min_chunk_duration_sec then drop the last chunk from the list
+        if (
+            len(chunks_to_extract_in_sec) > 0
+            and (chunks_to_extract_in_sec[-1][1] - chunks_to_extract_in_sec[-1][0]) < min_chunk_duration_sec
+        ):
+            return chunks_to_extract_in_sec[:-1]  # return all but the last chunk
+        return chunks_to_extract_in_sec
+
+    @classmethod
+    def input_schema(cls) -> dict[str, ts.ColumnType]:
+        return {
+            'audio': ts.AudioType(nullable=False),
+            'chunk_duration_sec': ts.FloatType(nullable=True),
+            'overlap_sec': ts.FloatType(nullable=True),
+            'min_chunk_duration_sec': ts.FloatType(nullable=True),
+        }
+
+    @classmethod
+    def output_schema(cls, *args: Any, **kwargs: Any) -> tuple[dict[str, ts.ColumnType], list[str]]:
+        param_names = ['chunk_duration_sec', 'min_chunk_duration_sec', 'overlap_sec']
+        params = dict(zip(param_names, args))
+        params.update(kwargs)
+
+        chunk_duration_sec = params['chunk_duration_sec']
+        min_chunk_duration_sec = params.get('min_chunk_duration_sec', 0.0)
+        overlap_sec = params.get('overlap_sec', 0.0)
+        if chunk_duration_sec <= 0.0:
+            raise excs.Error('chunk_duration_sec must be a positive number')
+        if chunk_duration_sec < min_chunk_duration_sec:
+            raise excs.Error('chunk_duration_sec must be at least min_chunk_duration_sec')
+        if overlap_sec >= chunk_duration_sec:
+            raise excs.Error('overlap_sec must be less than chunk_duration_sec')
+        return {
+            'start_time_sec': ts.FloatType(),
+            'end_time_sec': ts.FloatType(),
+            'audio_chunk': ts.AudioType(nullable=True),
+        }, []
+
+    def __next__(self) -> dict[str, Any]:
+        if self.next_pos >= len(self.chunks_to_extract_in_pts):
+            raise StopIteration
+        target_chunk_start, target_chunk_end = self.chunks_to_extract_in_pts[self.next_pos]
+        chunk_start_pts = 0
+        chunk_end_pts = 0
+        chunk_file = str(TempStore.create_path(extension=self.audio_path.suffix))
+        output_container = av.open(chunk_file, mode='w')
+        input_stream = self.container.streams.audio[0]
+        codec_name = AudioSplitter.__codec_map.get(input_stream.codec_context.name, input_stream.codec_context.name)
+        output_stream = output_container.add_stream(codec_name, rate=input_stream.codec_context.sample_rate)
+        assert isinstance(output_stream, av.audio.stream.AudioStream)
+        frame_count = 0
+        # Since frames don't align with chunk boundaries, we may have read an extra frame in previous iteration
+        # Seek to the nearest frame in stream at current chunk start time
+        self.container.seek(target_chunk_start, backward=True, stream=self.container.streams.audio[0])
+        while True:
+            try:
+                frame = next(self.container.decode(audio=0))
+            except EOFError as e:
+                raise excs.Error(f"Failed to read audio file '{self.audio_path}': {e}") from e
+            except StopIteration:
+                # no more frames to scan
+                break
+            if frame.pts < target_chunk_start:
+                # Current frame is behind chunk's start time, always get frame next to chunk's start time
+                continue
+            if frame.pts >= target_chunk_end:
+                # Frame has crossed the chunk boundary, it should be picked up by next chunk, throw away
+                # the current frame
+                break
+            frame_end = frame.pts + frame.samples
+            if frame_count == 0:
+                # Record start of the first frame
+                chunk_start_pts = frame.pts
+            # Write frame to output container
+            frame_count += 1
+            # If encode returns packets, write them to output container. Some encoders will buffer the frames.
+            output_container.mux(output_stream.encode(frame))
+            # record this frame's end as chunks end
+            chunk_end_pts = frame_end
+            # Check if frame's end has crossed the chunk boundary
+            if frame_end >= target_chunk_end:
+                break
+
+        # record result
+        if frame_count > 0:
+            # flush encoder
+            output_container.mux(output_stream.encode(None))
+            output_container.close()
+            result = {
+                'start_time_sec': round(float(chunk_start_pts * self.audio_time_base), 4),
+                'end_time_sec': round(float(chunk_end_pts * self.audio_time_base), 4),
+                'audio_chunk': chunk_file if frame_count > 0 else None,
+            }
+            _logger.debug('audio chunk result: %s', result)
+            self.next_pos += 1
+            return result
+        else:
+            # It's possible that there are no frames in the range of the last chunk, stop the iterator in this case.
+            # Note that start_time points at the first frame so case applies only for the last chunk
+            assert self.next_pos == len(self.chunks_to_extract_in_pts) - 1
+            self.next_pos += 1
+            raise StopIteration
+
+    def close(self) -> None:
+        self.container.close()
+
+    @classmethod
+    @deprecated('create() is deprecated; use `pixeltable.functions.audio.audio_splitter` instead', version='0.5.6')
+    def create(cls, **kwargs: Any) -> tuple[type[ComponentIterator], dict[str, Any]]:
+        return super()._create(**kwargs)

pixeltable/iterators/base.py

@@ -43,11 +43,17 @@ class ComponentIterator(ABC):
         """Close the iterator and release all resources"""
         raise NotImplementedError
 
-    @abstractmethod
-    def set_pos(self, pos: int) -> None:
+    def set_pos(self, pos: int, **kwargs: Any) -> None:
         """Set the iterator position to pos"""
-        raise NotImplementedError
+        pass
 
     @classmethod
     def create(cls, **kwargs: Any) -> tuple[type[ComponentIterator], dict[str, Any]]:
+        # TODO: This is still needed for compatibility with existing user-defined iterators; it will become deprecated
+        #     when the new decorator pattern is introduced for iterators
+        return cls._create(**kwargs)
+
+    @classmethod
+    def _create(cls, **kwargs: Any) -> tuple[type[ComponentIterator], dict[str, Any]]:
+        # create() variant that can be called by subclasses without generating a deprecation warning.
         return cls, kwargs