pixeltable 0.2.8__py3-none-any.whl → 0.2.10__py3-none-any.whl

pixeltable/functions/video.py

@@ -1,11 +1,16 @@
+import tempfile
 import uuid
+from pathlib import Path
 from typing import Optional
 
+import PIL.Image
 import av
+import numpy as np
 
 import pixeltable.env as env
 import pixeltable.func as func
 import pixeltable.type_system as ts
+from pixeltable.utils.code import local_public_names
 
 _format_defaults = {  # format -> (codec, ext)
     'wav': ('pcm_s16le', 'wav'),
@@ -30,6 +35,43 @@ _format_defaults = {  # format -> (codec, ext)
 #         output_container.mux(packet)
 
 
+@func.uda(
+    init_types=[ts.IntType()],
+    update_types=[ts.ImageType()],
+    value_type=ts.VideoType(),
+    requires_order_by=True,
+    allows_window=False,
+)
+class make_video(func.Aggregator):
+    def __init__(self, fps: int = 25):
+        """follows https://pyav.org/docs/develop/cookbook/numpy.html#generating-video"""
+        self.container: Optional[av.container.OutputContainer] = None
+        self.stream: Optional[av.stream.Stream] = None
+        self.fps = fps
+
+    def update(self, frame: PIL.Image.Image) -> None:
+        if frame is None:
+            return
+        if self.container is None:
+            (_, output_filename) = tempfile.mkstemp(suffix='.mp4', dir=str(env.Env.get().tmp_dir))
+            self.out_file = Path(output_filename)
+            self.container = av.open(str(self.out_file), mode='w')
+            self.stream = self.container.add_stream('h264', rate=self.fps)
+            self.stream.pix_fmt = 'yuv420p'
+            self.stream.width = frame.width
+            self.stream.height = frame.height
+
+        av_frame = av.VideoFrame.from_ndarray(np.array(frame.convert('RGB')), format='rgb24')
+        for packet in self.stream.encode(av_frame):
+            self.container.mux(packet)
+
+    def value(self) -> str:
+        for packet in self.stream.encode():
+            self.container.mux(packet)
+        self.container.close()
+        return str(self.out_file)
+
+
 _extract_audio_param_types = [
     ts.VideoType(nullable=False),
     ts.IntType(nullable=False),
@@ -75,26 +117,62 @@ def get_metadata(video: str) -> dict:
     """
     with av.open(video) as container:
         assert isinstance(container, av.container.InputContainer)
-        video_streams_info = [
-            {
-                'duration': stream.duration,
-                'frames': stream.frames,
-                'language': stream.language,
-                'average_rate': float(stream.average_rate) if stream.average_rate is not None else None,
-                'base_rate': float(stream.base_rate) if stream.base_rate is not None else None,
-                'guessed_rate': float(stream.guessed_rate) if stream.guessed_rate is not None else None,
-                'pix_fmt': getattr(stream.codec_context, 'pix_fmt', None),
-                'width': stream.width,
-                'height': stream.height,
-            }
-            for stream in container.streams
-            if isinstance(stream, av.video.stream.VideoStream)
-        ]
+        streams_info = [__get_stream_metadata(stream) for stream in container.streams]
         result = {
             'bit_exact': container.bit_exact,
             'bit_rate': container.bit_rate,
             'size': container.size,
             'metadata': container.metadata,
-            'streams': video_streams_info,  # TODO: Audio streams?
+            'streams': streams_info,
         }
     return result
+
+
+def __get_stream_metadata(stream: av.stream.Stream) -> dict:
+    if stream.type != 'audio' and stream.type != 'video':
+        return {'type': stream.type}  # Currently unsupported
+
+    codec_context = stream.codec_context
+    codec_context_md = {
+        'name': codec_context.name,
+        'codec_tag': codec_context.codec_tag.encode('unicode-escape').decode('utf-8'),
+        'profile': codec_context.profile,
+    }
+    metadata = {
+        'type': stream.type,
+        'duration': stream.duration,
+        'time_base': float(stream.time_base) if stream.time_base is not None else None,
+        'duration_seconds': float(stream.duration * stream.time_base)
+        if stream.duration is not None and stream.time_base is not None
+        else None,
+        'frames': stream.frames,
+        'metadata': stream.metadata,
+        'codec_context': codec_context_md,
+    }
+
+    if stream.type == 'audio':
+        # Additional metadata for audio
+        codec_context_md['channels'] = int(codec_context.channels) if codec_context.channels is not None else None
+    else:
+        assert stream.type == 'video'
+        # Additional metadata for video
+        codec_context_md['pix_fmt'] = getattr(stream.codec_context, 'pix_fmt', None)
+        metadata.update(
+            **{
+                'width': stream.width,
+                'height': stream.height,
+                'frames': stream.frames,
+                'average_rate': float(stream.average_rate) if stream.average_rate is not None else None,
+                'base_rate': float(stream.base_rate) if stream.base_rate is not None else None,
+                'guessed_rate': float(stream.guessed_rate) if stream.guessed_rate is not None else None,
+            }
+        )
+
+    return metadata
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/whisper.py

@@ -0,0 +1,74 @@
+from typing import TYPE_CHECKING, Optional
+
+import pixeltable as pxt
+
+if TYPE_CHECKING:
+    from whisper import Whisper
+
+
+@pxt.udf(
+    param_types=[
+        pxt.AudioType(),
+        pxt.StringType(),
+        pxt.JsonType(nullable=True),
+        pxt.FloatType(nullable=True),
+        pxt.FloatType(nullable=True),
+        pxt.FloatType(nullable=True),
+        pxt.BoolType(),
+        pxt.StringType(nullable=True),
+        pxt.BoolType(),
+        pxt.StringType(),
+        pxt.StringType(),
+        pxt.StringType(),
+        pxt.FloatType(nullable=True),
+        pxt.JsonType(nullable=True),
+    ]
+)
+def transcribe(
+    audio: str,
+    *,
+    model: str,
+    temperature: Optional[list[float]] = [0.0, 0.2, 0.4, 0.6, 0.8, 1.0],
+    compression_ratio_threshold: Optional[float] = 2.4,
+    logprob_threshold: Optional[float] = -1.0,
+    no_speech_threshold: Optional[float] = 0.6,
+    condition_on_previous_text: bool = True,
+    initial_prompt: Optional[str] = None,
+    word_timestamps: bool = False,
+    prepend_punctuations: str = '"\'“¿([{-',
+    append_punctuations: str = '"\'.。,，!！?？:：”)]}、',
+    decode_options: Optional[dict] = None,
+) -> dict:
+    import torch
+
+    if decode_options is None:
+        decode_options = {}
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    model = _lookup_model(model, device)
+    result = model.transcribe(
+        audio,
+        temperature=tuple(temperature),
+        compression_ratio_threshold=compression_ratio_threshold,
+        logprob_threshold=logprob_threshold,
+        no_speech_threshold=no_speech_threshold,
+        condition_on_previous_text=condition_on_previous_text,
+        initial_prompt=initial_prompt,
+        word_timestamps=word_timestamps,
+        prepend_punctuations=prepend_punctuations,
+        append_punctuations=append_punctuations,
+        **decode_options,
+    )
+    return result
+
+
+def _lookup_model(model_id: str, device: str) -> 'Whisper':
+    import whisper
+
+    key = (model_id, device)
+    if key not in _model_cache:
+        model = whisper.load_model(model_id, device)
+        _model_cache[key] = model
+    return _model_cache[key]
+
+
+_model_cache: dict[tuple[str, str], 'Whisper'] = {}

pixeltable/globals.py

@@ -213,7 +213,7 @@ def move(path: str, new_path: str) -> None:
     obj = Catalog.get().paths[p]
     Catalog.get().paths.move(p, new_p)
     new_dir = Catalog.get().paths[new_p.parent]
-    obj.move(new_p.name, new_dir._id)
+    obj._move(new_p.name, new_dir._id)
 
 
 def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> None:

pixeltable/io/__init__.py

@@ -1,4 +1,14 @@
+from .external_store import ExternalStore, SyncStatus
 from .globals import create_label_studio_project
 from .hf_datasets import import_huggingface_dataset
 from .pandas import import_csv, import_excel, import_pandas
 from .parquet import import_parquet
+
+
+__default_dir = set(symbol for symbol in dir() if not symbol.startswith('_'))
+__removed_symbols = {'globals', 'hf_datasets', 'pandas', 'parquet'}
+__all__ = sorted(list(__default_dir - __removed_symbols))
+
+
+def __dir__():
+    return __all__

pixeltable/io/external_store.py

@@ -0,0 +1,370 @@
+from __future__ import annotations
+
+import abc
+import itertools
+import logging
+import time
+from dataclasses import dataclass
+from typing import Any, Optional
+from uuid import UUID
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+from pixeltable import Table, Column
+import sqlalchemy as sql
+
+from pixeltable.catalog import TableVersion
+
+_logger = logging.getLogger('pixeltable')
+
+
+class ExternalStore(abc.ABC):
+    """
+    Abstract base class that represents an external data store that is linked to a Pixeltable
+    table. Subclasses of `ExternalStore` provide functionality for synchronizing between Pixeltable
+    and stateful external stores.
+    """
+
+    def __init__(self, name: str) -> None:
+        self.__name = name
+
+    @property
+    def name(self) -> str:
+        return self.__name
+
+    @abc.abstractmethod
+    def link(self, tbl_version: TableVersion, conn: sql.Connection) -> None:
+        """
+        Called by `TableVersion.link()` to implement store-specific logic.
+        """
+
+    @abc.abstractmethod
+    def unlink(self, tbl_version: TableVersion, conn: sql.Connection) -> None:
+        """
+        Called by `TableVersion.unlink()` to implement store-specific logic.
+        """
+
+    @abc.abstractmethod
+    def get_local_columns(self) -> list[Column]:
+        """
+        Gets a list of all local (Pixeltable) columns that are associated with this external store.
+        """
+
+    @abc.abstractmethod
+    def sync(self, t: Table, export_data: bool, import_data: bool) -> SyncStatus:
+        """
+        Called by `Table.sync()` to implement store-specific synchronization logic.
+        """
+
+    @abc.abstractmethod
+    def as_dict(self) -> dict[str, Any]: ...
+
+    @classmethod
+    @abc.abstractmethod
+    def from_dict(cls, md: dict[str, Any]) -> ExternalStore: ...
+
+
+class Project(ExternalStore, abc.ABC):
+    """
+    An `ExternalStore` that represents a labeling project. Extends `ExternalStore` with a few
+    additional capabilities specific to such projects.
+    """
+    def __init__(self, name: str, col_mapping: dict[Column, str], stored_proxies: Optional[dict[Column, Column]]):
+        super().__init__(name)
+        self._col_mapping = col_mapping
+
+        # A mapping from original columns to proxy columns. A proxy column is an identical copy of a column that is
+        # guaranteed to be stored; the Project will dynamically create and tear down proxy columns as needed. There
+        # are two reasons this might happen:
+        # (i) to force computed media data to be persisted; or
+        # (ii) to force media data to be materialized in a particular location.
+        # For each entry (k, v) in the dict, `v` is the stored proxy column for `k`. The proxy column `v` will
+        # necessarily be a column of the table to which this project is linked, but `k` need not be; it might be a
+        # column of a base table.
+        # Note from aaron-siegel: This methodology is inefficient in the case where a table has many views with a high
+        # proportion of overlapping rows, all proxying the same base column.
+        if stored_proxies is None:
+            self.stored_proxies: dict[Column, Column] = {}
+        else:
+            self.stored_proxies = stored_proxies
+
+    def get_local_columns(self) -> list[Column]:
+        return list(self.col_mapping.keys())
+
+    def link(self, tbl_version: TableVersion, conn: sql.Connection) -> None:
+        # All of the media columns being linked need to either be stored computed columns, or else have stored proxies.
+        # This ensures that the media in those columns resides in the media store.
+        # First determine which columns (if any) need stored proxies, but don't have one yet.
+        stored_proxies_needed: list[Column] = []
+        for col in self.col_mapping.keys():
+            if col.col_type.is_media_type() and not (col.is_stored and col.is_computed):
+                # If this column is already proxied in some other Project, use the existing proxy to avoid
+                # duplication. Otherwise, we'll create a new one.
+                for store in tbl_version.external_stores.values():
+                    if isinstance(store, Project) and col in store.stored_proxies:
+                        self.stored_proxies[col] = store.stored_proxies[col]
+                        break
+                if col not in self.stored_proxies:
+                    # We didn't find it in an existing Project
+                    stored_proxies_needed.append(col)
+        if len(stored_proxies_needed) > 0:
+            _logger.info(f'Creating stored proxies for columns: {[col.name for col in stored_proxies_needed]}')
+            # Create stored proxies for columns that need one. Increment the schema version
+            # accordingly.
+            tbl_version.version += 1
+            preceding_schema_version = tbl_version.schema_version
+            tbl_version.schema_version = tbl_version.version
+            proxy_cols = [self.create_stored_proxy(tbl_version, col) for col in stored_proxies_needed]
+            # Add the columns; this will also update table metadata.
+            tbl_version._add_columns(proxy_cols, conn)
+            # We don't need to retain `UpdateStatus` since the stored proxies are intended to be
+            # invisible to the user.
+            tbl_version._update_md(time.time(), conn, preceding_schema_version=preceding_schema_version)
+
+    def unlink(self, tbl_version: TableVersion, conn: sql.Connection) -> None:
+        # Determine which stored proxies can be deleted. (A stored proxy can be deleted if it is not referenced by
+        # any *other* external store for this table.)
+        deletions_needed: set[Column] = set(self.stored_proxies.values())
+        for name, store in tbl_version.external_stores.items():
+            if name != self.name:
+                deletions_needed = deletions_needed.difference(set(store.stored_proxies.values()))
+        if len(deletions_needed) > 0:
+            _logger.info(f'Removing stored proxies for columns: {[col.name for col in deletions_needed]}')
+            # Delete stored proxies that are no longer needed.
+            tbl_version.version += 1
+            preceding_schema_version = tbl_version.schema_version
+            tbl_version.schema_version = tbl_version.version
+            tbl_version._drop_columns(deletions_needed)
+            self.stored_proxies.clear()
+            tbl_version._update_md(time.time(), conn, preceding_schema_version=preceding_schema_version)
+
+    def create_stored_proxy(self, tbl_version: TableVersion, col: Column) -> Column:
+        """
+        Creates a proxy column for the specified column. The proxy column will be created in the specified
+        `TableVersion`.
+        """
+        from pixeltable import exprs
+
+        assert col.col_type.is_media_type() and not (col.is_stored and col.is_computed) and col not in self.stored_proxies
+        proxy_col = Column(
+            name=None,
+            # Force images in the proxy column to be materialized inside the media store, in a normalized format.
+            # TODO(aaron-siegel): This is a temporary solution and it will be replaced by a proper `destination`
+            #   parameter for computed columns. Among other things, this solution does not work for video or audio.
+            #   Once `destination` is implemented, it can be replaced with a simple `ColumnRef`.
+            computed_with=exprs.ColumnRef(col).apply(lambda x: x, col_type=col.col_type),
+            stored=True,
+            col_id=tbl_version.next_col_id,
+            sa_col_type=col.col_type.to_sa_type(),
+            schema_version_add=tbl_version.schema_version
+        )
+        proxy_col.tbl = tbl_version
+        tbl_version.next_col_id += 1
+        self.stored_proxies[col] = proxy_col
+        return proxy_col
+
+    @property
+    def col_mapping(self) -> dict[Column, str]:
+        return self._col_mapping
+
+    @abc.abstractmethod
+    def get_export_columns(self) -> dict[str, ts.ColumnType]:
+        """
+        Returns the names and Pixeltable types that this `Project` expects to see in a data export. The keys
+        of the `dict` are the names of data fields in the external store, not Pixeltable columns.
+
+        Returns:
+            A `dict` mapping names of external data fields to their expected Pixeltable types.
+        """
+
+    @abc.abstractmethod
+    def get_import_columns(self) -> dict[str, ts.ColumnType]:
+        """
+        Returns the names and Pixeltable types that this `Project` provides in a data import.
+
+        Returns:
+            A `dict` mapping names of provided columns to their Pixeltable types.
+        """
+
+    @abc.abstractmethod
+    def delete(self) -> None:
+        """
+        Deletes this `Project` and all associated (externally stored) data.
+        """
+
+    @classmethod
+    def validate_columns(
+            cls,
+            table: Table,
+            export_cols: dict[str, ts.ColumnType],
+            import_cols: dict[str, ts.ColumnType],
+            col_mapping: Optional[dict[str, str]]
+    ) -> dict[Column, str]:
+        """
+        Verifies that the specified `col_mapping` is valid. In particular, checks that:
+        (i) the keys of `col_mapping` are valid columns of the specified `Table`;
+        (ii) the values of `col_mapping` are valid external columns (i.e., they appear in either `export_cols` or
+            `import_cols`; and
+        (iii) the Pixeltable types of the `col_mapping` keys are consistent with the expected types of the corresponding
+            external (import or export) columns.
+        If validation fails, an exception will be raised. If validation succeeds, a new mapping will be returned
+        in which the Pixeltable column names are resolved to the corresponding `Column` objects.
+        """
+        is_user_specified_col_mapping = col_mapping is not None
+        if col_mapping is None:
+            col_mapping = {col: col for col in itertools.chain(export_cols.keys(), import_cols.keys())}
+
+        resolved_col_mapping: dict[Column, str] = {}
+
+        # Validate names
+        t_cols = table.column_names()
+        for t_col, ext_col in col_mapping.items():
+            if t_col not in t_cols:
+                if is_user_specified_col_mapping:
+                    raise excs.Error(
+                        f'Column name `{t_col}` appears as a key in `col_mapping`, but Table `{table.get_name()}` '
+                        'contains no such column.'
+                    )
+                else:
+                    raise excs.Error(
+                        f'Column `{t_col}` does not exist in Table `{table.get_name()}`. Either add a column `{t_col}`, '
+                        f'or specify a `col_mapping` to associate a different column with the external field `{ext_col}`.'
+                    )
+            if ext_col not in export_cols and ext_col not in import_cols:
+                raise excs.Error(
+                    f'Column name `{ext_col}` appears as a value in `col_mapping`, but the external store '
+                    f'configuration has no column `{ext_col}`.'
+                )
+            col = table[t_col].col
+            resolved_col_mapping[col] = ext_col
+        # Validate column specs
+        t_col_types = table.column_types()
+        for t_col, ext_col in col_mapping.items():
+            t_col_type = t_col_types[t_col]
+            if ext_col in export_cols:
+                # Validate that the table column can be assigned to the external column
+                ext_col_type = export_cols[ext_col]
+                if not ext_col_type.is_supertype_of(t_col_type):
+                    raise excs.Error(
+                        f'Column `{t_col}` cannot be exported to external column `{ext_col}` (incompatible types; expecting `{ext_col_type}`)'
+                    )
+            if ext_col in import_cols:
+                # Validate that the external column can be assigned to the table column
+                if table._tbl_version_path.get_column(t_col).is_computed:
+                    raise excs.Error(
+                        f'Column `{t_col}` is a computed column, which cannot be populated from an external column'
+                    )
+                ext_col_type = import_cols[ext_col]
+                if not t_col_type.is_supertype_of(ext_col_type):
+                    raise excs.Error(
+                        f'Column `{t_col}` cannot be imported from external column `{ext_col}` (incompatible types; expecting `{ext_col_type}`)'
+                    )
+        return resolved_col_mapping
+
+    @classmethod
+    def _column_as_dict(cls, col: Column) -> dict[str, Any]:
+        return {'tbl_id': str(col.tbl.id), 'col_id': col.id}
+
+    @classmethod
+    def _column_from_dict(cls, d: dict[str, Any]) -> Column:
+        from pixeltable.catalog import Catalog
+
+        tbl_id = UUID(d['tbl_id'])
+        col_id = d['col_id']
+        return Catalog.get().tbl_versions[(tbl_id, None)].cols_by_id[col_id]
+
+
+@dataclass(frozen=True)
+class SyncStatus:
+    external_rows_created: int = 0
+    external_rows_deleted: int = 0
+    external_rows_updated: int = 0
+    pxt_rows_updated: int = 0
+    num_excs: int = 0
+
+    def combine(self, other: 'SyncStatus') -> 'SyncStatus':
+        return SyncStatus(
+            external_rows_created=self.external_rows_created + other.external_rows_created,
+            external_rows_deleted=self.external_rows_deleted + other.external_rows_deleted,
+            external_rows_updated=self.external_rows_updated + other.external_rows_updated,
+            pxt_rows_updated=self.pxt_rows_updated + other.pxt_rows_updated,
+            num_excs=self.num_excs + other.num_excs
+        )
+
+    @classmethod
+    def empty(cls) -> 'SyncStatus':
+        return SyncStatus(0, 0, 0, 0, 0)
+
+
+class MockProject(Project):
+    """A project that cannot be synced, used mainly for testing."""
+    def __init__(
+            self,
+            name: str,
+            export_cols: dict[str, ts.ColumnType],
+            import_cols: dict[str, ts.ColumnType],
+            col_mapping: dict[Column, str],
+            stored_proxies: Optional[dict[Column, Column]] = None
+    ):
+        super().__init__(name, col_mapping, stored_proxies)
+        self.export_cols = export_cols
+        self.import_cols = import_cols
+        self.__is_deleted = False
+
+    @classmethod
+    def create(
+            cls,
+            t: Table,
+            name: str,
+            export_cols: dict[str, ts.ColumnType],
+            import_cols: dict[str, ts.ColumnType],
+            col_mapping: Optional[dict[str, str]] = None
+    ) -> 'MockProject':
+        col_mapping = cls.validate_columns(t, export_cols, import_cols, col_mapping)
+        return cls(name, export_cols, import_cols, col_mapping)
+
+    def get_export_columns(self) -> dict[str, ts.ColumnType]:
+        return self.export_cols
+
+    def get_import_columns(self) -> dict[str, ts.ColumnType]:
+        return self.import_cols
+
+    def sync(self, t: Table, export_data: bool, import_data: bool) -> NotImplemented:
+        raise NotImplementedError()
+
+    def delete(self) -> None:
+        self.__is_deleted = True
+
+    @property
+    def is_deleted(self) -> bool:
+        return self.__is_deleted
+
+    def as_dict(self) -> dict[str, Any]:
+        return {
+            'name': self.name,
+            'export_cols': {k: v.as_dict() for k, v in self.export_cols.items()},
+            'import_cols': {k: v.as_dict() for k, v in self.import_cols.items()},
+            'col_mapping': [[self._column_as_dict(k), v] for k, v in self.col_mapping.items()],
+            'stored_proxies': [[self._column_as_dict(k), self._column_as_dict(v)] for k, v in self.stored_proxies.items()]
+        }
+
+    @classmethod
+    def from_dict(cls, md: dict[str, Any]) -> MockProject:
+        return cls(
+            md['name'],
+            {k: ts.ColumnType.from_dict(v) for k, v in md['export_cols'].items()},
+            {k: ts.ColumnType.from_dict(v) for k, v in md['import_cols'].items()},
+            {cls._column_from_dict(entry[0]): entry[1] for entry in md['col_mapping']},
+            {cls._column_from_dict(entry[0]): cls._column_from_dict(entry[1]) for entry in md['stored_proxies']}
+        )
+
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, MockProject):
+            return False
+        return self.name == other.name
+
+    def __hash__(self) -> int:
+        return hash(self.name)
+
+    def __repr__(self) -> str:
+        return f'MockProject `{self.name}`'