pixeltable 0.2.21__py3-none-any.whl → 0.2.22__py3-none-any.whl

pixeltable/functions/replicate.py

@@ -0,0 +1,72 @@
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs)
+that wrap various endpoints from the Replicate API. In order to use them, you must
+first `pip install replicate` and configure your Replicate credentials, as described in
+the [Working with Replicate](https://pixeltable.readme.io/docs/working-with-replicate) tutorial.
+"""
+
+from typing import TYPE_CHECKING, Any
+
+import pixeltable as pxt
+from pixeltable.env import Env, register_client
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import replicate  # type: ignore[import-untyped]
+
+
+@register_client('replicate')
+def _(api_token: str) -> 'replicate.Client':
+    import replicate
+    return replicate.Client(api_token=api_token)
+
+
+def _replicate_client() -> 'replicate.Client':
+    return Env.get().get_client('replicate')
+
+
+@pxt.udf
+def run(
+    input: dict[str, Any],
+    *,
+    ref: str,
+) -> dict[str, Any]:
+    """
+    Run a model on Replicate.
+
+    For additional details, see: <https://replicate.com/docs/topics/models/run-a-model>
+
+    __Requirements:__
+
+    - `pip install replicate`
+
+    Args:
+        input: The input parameters for the model.
+        ref: The name of the model to run.
+
+    Returns:
+        The output of the model.
+
+    Examples:
+        Add a computed column that applies the model `meta/meta-llama-3-8b-instruct`
+        to an existing Pixeltable column `tbl.prompt` of the table `tbl`:
+
+        >>> input = {'system_prompt': 'You are a helpful assistant.', 'prompt': tbl.prompt}
+        ... tbl['response'] = run(input, ref='meta/meta-llama-3-8b-instruct')
+
+        Add a computed column that uses the model `black-forest-labs/flux-schnell`
+        to generate images from an existing Pixeltable column `tbl.prompt`:
+
+        >>> input = {'prompt': tbl.prompt, 'go_fast': True, 'megapixels': '1'}
+        ... tbl['response'] = run(input, ref='black-forest-labs/flux-schnell')
+        ... tbl['image'] = tbl.response.output[0].astype(pxt.Image)
+    """
+    Env.get().require_package('replicate')
+    return _replicate_client().run(ref, input, use_file_output=False)
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/string.py

@@ -283,6 +283,15 @@ def isspace(self: str) -> bool:
     """
     return self.isspace()
 
+@pxt.udf
+def join(sep: str, elements: list) -> str:
+    """
+    Return a string which is the concatenation of the strings in `elements`.
+
+    Equivalent to [`str.join()`](https://docs.python.org/3/library/stdtypes.html#str.join)
+    """
+    return sep.join(elements)
+
 @pxt.udf(is_method=True)
 def len(self: str) -> int:
     """

pixeltable/globals.py

@@ -1,6 +1,6 @@
 import dataclasses
 import logging
-from typing import Any, Iterable, Optional, Union
+from typing import Any, Iterable, Optional, Union, Literal
 from uuid import UUID
 
 import pandas as pd
@@ -33,6 +33,7 @@ def create_table(
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
+    media_validation: Literal['on_read', 'on_write'] = 'on_write'
 ) -> catalog.Table:
     """Create a new base table.
 
@@ -44,6 +45,9 @@ def create_table(
             table.
         num_retained_versions: Number of versions of the table to retain.
         comment: An optional comment; its meaning is user-defined.
+        media_validation: Media validation policy for the table.
+            - `'on_read'`: validate media files at query time
+            - `'on_write'`: validate media files during insert/update operations
 
     Returns:
         A handle to the newly created [`Table`][pixeltable.Table].
@@ -89,14 +93,8 @@ def create_table(
             raise excs.Error('primary_key must be a single column name or a list of column names')
 
     tbl = catalog.InsertableTable._create(
-        dir._id,
-        path.name,
-        schema,
-        df,
-        primary_key=primary_key,
-        num_retained_versions=num_retained_versions,
-        comment=comment,
-    )
+        dir._id, path.name, schema, df, primary_key=primary_key, num_retained_versions=num_retained_versions,
+        comment=comment, media_validation=catalog.MediaValidation.validated(media_validation, 'media_validation'))
     Catalog.get().paths[path] = tbl
 
     _logger.info(f'Created table `{path_str}`.')
@@ -112,6 +110,7 @@ def create_view(
     iterator: Optional[tuple[type[ComponentIterator], dict[str, Any]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
+    media_validation: Literal['on_read', 'on_write'] = 'on_write',
     ignore_errors: bool = False,
 ) -> Optional[catalog.Table]:
     """Create a view of an existing table object (which itself can be a view or a snapshot or a base table).
@@ -177,17 +176,10 @@ def create_view(
         iterator_class, iterator_args = iterator
 
     view = catalog.View._create(
-        dir._id,
-        path.name,
-        base=tbl_version_path,
-        additional_columns=additional_columns,
-        predicate=where,
-        is_snapshot=is_snapshot,
-        iterator_cls=iterator_class,
-        iterator_args=iterator_args,
-        num_retained_versions=num_retained_versions,
-        comment=comment,
-    )
+        dir._id, path.name, base=tbl_version_path, additional_columns=additional_columns, predicate=where,
+        is_snapshot=is_snapshot, iterator_cls=iterator_class, iterator_args=iterator_args,
+        num_retained_versions=num_retained_versions, comment=comment,
+        media_validation=catalog.MediaValidation.validated(media_validation, 'media_validation'))
     Catalog.get().paths[path] = view
     _logger.info(f'Created view `{path_str}`.')
     FileCache.get().emit_eviction_warnings()

pixeltable/index/btree.py

@@ -1,13 +1,16 @@
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
 import sqlalchemy as sql
 
+# TODO: why does this import result in a circular import, but the one im embedding_index.py doesn't?
+# import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 from pixeltable import catalog, exprs
 from pixeltable.func.udf import udf
-
 from .base import IndexBase
 
+if TYPE_CHECKING:
+    import pixeltable.exprs
 
 class BtreeIndex(IndexBase):
     """
@@ -15,6 +18,8 @@ class BtreeIndex(IndexBase):
     """
     MAX_STRING_LEN = 256
 
+    value_expr: 'pixeltable.exprs.Expr'
+
     @staticmethod
     @udf
     def str_filter(s: Optional[str]) -> Optional[str]:
@@ -25,7 +30,14 @@ class BtreeIndex(IndexBase):
     def __init__(self, c: 'catalog.Column'):
         if not c.col_type.is_scalar_type() and not c.col_type.is_media_type():
             raise excs.Error(f'Index on column {c.name}: B-tree index requires scalar or media type, got {c.col_type}')
-        self.value_expr = BtreeIndex.str_filter(exprs.ColumnRef(c)) if c.col_type.is_string_type() else exprs.ColumnRef(c)
+        if c.col_type.is_media_type():
+            # an index on a media column is an index on the file url
+            # no validation for media columns: we're only interested in the string value
+            self.value_expr = exprs.ColumnRef(c, perform_validation=False)
+        else:
+            self.value_expr = (
+                BtreeIndex.str_filter(exprs.ColumnRef(c)) if c.col_type.is_string_type() else exprs.ColumnRef(c)
+            )
 
     def index_value_expr(self) -> 'exprs.Expr':
         return self.value_expr
@@ -52,3 +64,4 @@ class BtreeIndex(IndexBase):
     @classmethod
     def from_dict(cls, c: 'catalog.Column', d: dict) -> 'BtreeIndex':
         return cls(c)
+

pixeltable/index/embedding_index.py

@@ -86,8 +86,8 @@ class EmbeddingIndex(IndexBase):
         )
         idx.create(bind=conn)
 
-    def similarity_clause(self, val_column: catalog.Column, item: Any) -> sql.ClauseElement:
-        """Create a ClauseElement that represents '<val_column> <op> <item>'"""
+    def similarity_clause(self, val_column: catalog.Column, item: Any) -> sql.ColumnElement:
+        """Create a ColumnElement that represents '<val_column> <op> <item>'"""
         assert isinstance(item, (str, PIL.Image.Image))
         if isinstance(item, str):
             assert self.string_embed is not None
@@ -104,8 +104,8 @@ class EmbeddingIndex(IndexBase):
             assert self.metric == self.Metric.L2
             return val_column.sa_col.l2_distance(embedding)
 
-    def order_by_clause(self, val_column: catalog.Column, item: Any, is_asc: bool) -> sql.ClauseElement:
-        """Create a ClauseElement that is used in an ORDER BY clause"""
+    def order_by_clause(self, val_column: catalog.Column, item: Any, is_asc: bool) -> sql.ColumnElement:
+        """Create a ColumnElement that is used in an ORDER BY clause"""
         assert isinstance(item, (str, PIL.Image.Image))
         embedding: Optional[np.ndarray] = None
         if isinstance(item, str):

pixeltable/io/__init__.py

@@ -1,10 +1,9 @@
 from .external_store import ExternalStore, SyncStatus
-from .globals import create_label_studio_project, import_rows, import_json
+from .globals import create_label_studio_project, export_images_as_fo_dataset, import_json, import_rows
 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))

pixeltable/io/fiftyone.py

@@ -0,0 +1,178 @@
+import os
+from typing import Iterator, Optional, Union
+
+import fiftyone as fo  # type: ignore[import-untyped]
+import fiftyone.utils.data as foud  # type: ignore[import-untyped]
+import PIL.Image
+import puremagic
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+from pixeltable import exprs
+from pixeltable.env import Env
+
+
+class PxtImageDatasetImporter(foud.LabeledImageDatasetImporter):
+    """
+    Implementation of a FiftyOne `DatasetImporter` that reads image data from a Pixeltable table.
+    """
+    __image_format: str  # format to use for any exported images that are not already stored on disk
+    __labels: dict[str, tuple[exprs.Expr, type[fo.Label]]]  # label_name -> (expr, label_cls)
+    __image_idx: int  # index of the image expr in the select list
+    __localpath_idx: Optional[int]  # index of the image localpath in the select list, if present
+    __row_iter: Iterator[list]  # iterator over the table rows, to be convered to FiftyOne samples
+
+    def __init__(
+        self,
+        tbl: pxt.Table,
+        image: exprs.Expr,
+        image_format: str,
+        classifications: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
+        detections: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
+        dataset_dir: Optional[os.PathLike] = None,
+        shuffle: bool = False,
+        seed: Union[int, float, str, bytes, bytearray, None] = None,
+        max_samples: Optional[int] = None,
+    ):
+        super().__init__(
+            dataset_dir=dataset_dir,
+            shuffle=shuffle,
+            seed=seed,
+            max_samples=max_samples
+        )
+
+        self.__image_format = image_format
+
+        label_categories = [
+            (classifications, fo.Classifications, 'classifications'),
+            (detections, fo.Detections, 'detections'),
+        ]
+
+        # Construct the labels. First add labels for all label types that have named dictionaries.
+        self.__labels = {}
+        for exprs_, label_cls, _ in label_categories:
+            if isinstance(exprs_, dict):
+                for label_name, expr in exprs_.items():
+                    if not label_name.isidentifier():
+                        raise excs.Error(f"Invalid label name: {label_name}")
+                    if label_name in self.__labels:
+                        raise excs.Error(f"Duplicate label name: {label_name}")
+                    self.__labels[label_name] = (expr, label_cls)
+
+        # Now add the remaining labels, assigning unused default names.
+        for exprs_, label_cls, default_name in label_categories:
+            if exprs_ is None or isinstance(exprs_, dict):
+                continue
+            if isinstance(exprs_, exprs.Expr):
+                exprs_ = [exprs_]
+            assert isinstance(exprs_, list)
+            for expr in exprs_:
+                if default_name not in self.__labels:
+                    name = default_name
+                else:
+                    i = 1
+                    while f'{default_name}_{i}' in self.__labels:
+                        i += 1
+                    name = f'{default_name}_{i}'
+                self.__labels[name] = (expr, label_cls)
+
+        # Build the select list:
+        # - Labels first, in the order they appear in self.__labels
+        # - Then the `image` expr
+        # - Then `image.localpath`, if `images` is a stored columnref
+
+        selection = [expr for expr, _ in self.__labels.values()]
+        self.__image_idx = len(selection)
+        selection.append(image)
+
+        if isinstance(image, exprs.ColumnRef) and image.col.is_stored:
+            # A stored image column; we can use the existing localpaths
+            self.__localpath_idx = len(selection)
+            selection.append(image.localpath)
+        else:
+            self.__localpath_idx = None
+
+        df = tbl.select(*selection)
+        self.__row_iter = df._output_row_iterator()
+
+    def __next__(self) -> tuple[str, Optional[fo.ImageMetadata], Optional[dict[str, fo.Label]]]:
+        row = next(self.__row_iter)
+        img = row[self.__image_idx]
+        assert isinstance(img, PIL.Image.Image)
+        if self.__localpath_idx is not None:
+            # Use the existing localpath of the stored image
+            file = row[self.__localpath_idx]
+            assert isinstance(file, str)
+        else:
+            # Write the dynamically created image to a temp file
+            file = str(Env.get().create_tmp_path(f'.{self.__image_format}'))
+            img.save(file, format=self.__image_format)
+
+        metadata = fo.ImageMetadata(
+            size_bytes=os.path.getsize(file),
+            mime_type=puremagic.from_file(file, mime=True),
+            width=img.width,
+            height=img.height,
+            filepath=file,
+            num_channels=len(img.getbands()),
+        )
+
+        labels: dict[str, fo.Label] = {}
+        for idx, (label_name, (_, label_cls)) in enumerate(self.__labels.items()):
+            label_data = row[idx]
+            if label_data is None:
+                continue
+
+            label: fo.Label
+            if label_cls is fo.Classifications:
+                label = fo.Classifications(classifications=self.__as_fo_classifications(label_data))
+            elif label_cls is fo.Detections:
+                label = fo.Detections(detections=self.__as_fo_detections(label_data))
+            else:
+                assert False
+            labels[label_name] = label
+
+        return file, metadata, labels
+
+    def __as_fo_classifications(self, data: list) -> list[fo.Classification]:
+        if not isinstance(data, list) or any('label' not in entry for entry in data):
+            raise excs.Error(
+                f'Invalid classifications data: {data}\n'
+                "(Expected a list of dicts, each containing a 'label' key)"
+            )
+        return [
+            fo.Classification(label=entry['label'], confidence=entry.get('confidence'))
+            for entry in data
+        ]
+
+    def __as_fo_detections(self, data: list) -> list[fo.Detections]:
+        if not isinstance(data, list) or any('label' not in entry or 'bounding_box' not in entry for entry in data):
+            raise excs.Error(
+                f'Invalid detections data: {data}\n'
+                "(Expected a list of dicts, each containing a 'label' and 'bounding_box' key)"
+            )
+        return [
+            fo.Detection(label=entry['label'], bounding_box=entry['bounding_box'], confidence=entry.get('confidence'))
+            for entry in data
+        ]
+
+    @property
+    def has_dataset_info(self) -> bool:
+        return False
+
+    @property
+    def has_image_metadata(self) -> bool:
+        return True
+
+    @property
+    def label_cls(self) -> dict[str, type]:
+        return {label_name: label_cls for label_name, (_, label_cls) in self.__labels.items()}
+
+    def setup(self) -> None:
+        pass
+
+    def get_dataset_info(self) -> dict:
+        pass
+
+    def close(self, *args) -> None:
+        pass

pixeltable/io/globals.py

@@ -1,10 +1,14 @@
-from typing import Any, Literal, Optional, Union
+from typing import TYPE_CHECKING, Any, Literal, Optional, Union
 
 import pixeltable as pxt
 import pixeltable.exceptions as excs
-from pixeltable import Table
+from pixeltable import Table, exprs
+from pixeltable.env import Env
 from pixeltable.io.external_store import SyncStatus
 
+if TYPE_CHECKING:
+    import fiftyone as fo  # type: ignore[import-untyped]
+
 
 def create_label_studio_project(
         t: Table,
@@ -116,6 +120,8 @@ def create_label_studio_project(
                 s3_configuration={'bucket': 'my-bucket', 'region_name': 'us-east-2'}
             )
     """
+    Env.get().require_package('label_studio_sdk')
+
     from pixeltable.io.label_studio import LabelStudioProject
 
     ls_project = LabelStudioProject.create(
@@ -267,3 +273,91 @@ def import_json(
         contents = urllib.request.urlopen(filepath_or_url).read()
     data = json.loads(contents, **kwargs)
     return import_rows(tbl_path, data, schema_overrides=schema_overrides, primary_key=primary_key, num_retained_versions=num_retained_versions, comment=comment)
+
+
+def export_images_as_fo_dataset(
+    tbl: pxt.Table,
+    images: exprs.Expr,
+    image_format: str = 'webp',
+    classifications: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
+    detections: Union[exprs.Expr, list[exprs.Expr], dict[str, exprs.Expr], None] = None,
+) -> 'fo.Dataset':
+    """
+    Export images from a Pixeltable table as a Voxel51 dataset. The data must consist of a single column
+    (or expression) containing image data, along with optional additional columns containing labels. Currently, only
+    classification and detection labels are supported.
+
+    The [Working with Voxel51 in Pixeltable](https://docs.pixeltable.com/docs/working-with-voxel51) tutorial contains a
+    fully worked example showing how to export data from a Pixeltable table and load it into Voxel51.
+
+    Images in the dataset that already exist on disk will be exported directly, in whatever format they
+    are stored in. Images that are not already on disk (such as frames extracted using a
+    [`FrameIterator`][pixeltable.iterators.FrameIterator]) will first be written to disk in the specified
+    `image_format`.
+
+    The label parameters accept one or more sets of labels of each type. If a single `Expr` is provided, then it will
+    be exported as a single set of labels with a default name such as `classifications`.
+    (The single set of labels may still containing multiple individual labels; see below.)
+    If a list of `Expr`s is provided, then each one will be exported as a separate set of labels with a default name
+    such as `classifications`, `classifications_1`, etc. If a dictionary of `Expr`s is provided, then each entry will
+    be exported as a set of labels with the specified name.
+
+    __Requirements:__
+
+    - `pip install fiftyone`
+
+    Args:
+        tbl: The table from which to export data.
+        images: A column or expression that contains the images to export.
+        image_format: The format to use when writing out images for export.
+        classifications: Optional image classification labels. If a single `Expr` is provided, it must be a table
+            column or an expression that evaluates to a list of dictionaries. Each dictionary in the list corresponds
+            to an image class and must have the following structure:
+
+            ```python
+            {'label': 'zebra', 'confidence': 0.325}
+            ```
+
+            If multiple `Expr`s are provided, each one must evaluate to a list of such dictionaries.
+        detections: Optional image detection labels. If a single `Expr` is provided, it must be a table column or an
+            expression that evaluates to a list of dictionaries. Each dictionary in the list corresponds to an image
+            detection, and must have the following structure:
+
+            ```python
+            {
+                'label': 'giraffe',
+                'confidence': 0.99,
+                'bounding_box': [0.081, 0.836, 0.202, 0.136]  # [x, y, w, h], fractional coordinates
+            }
+            ```
+
+            If multiple `Expr`s are provided, each one must evaluate to a list of such dictionaries.
+
+    Returns:
+        A Voxel51 dataset.
+
+    Example:
+        Export the images in the `image` column of the table `tbl` as a Voxel51 dataset, using classification
+        labels from `tbl.classifications`:
+
+        >>> export_as_fiftyone(
+        ...     tbl,
+        ...     tbl.image,
+        ...     classifications=tbl.classifications
+        ... )
+
+        See the [Working with Voxel51 in Pixeltable](https://docs.pixeltable.com/docs/working-with-voxel51) tutorial
+        for a fully worked example.
+    """
+    Env.get().require_package('fiftyone')
+
+    import fiftyone as fo
+
+    from pixeltable.io.fiftyone import PxtImageDatasetImporter
+
+    if not images.col_type.is_image_type():
+        raise excs.Error(f'`images` must be an expression of type Image (got {images.col_type._to_base_str()})')
+
+    return fo.Dataset.from_importer(PxtImageDatasetImporter(
+        tbl, images, image_format, classifications=classifications, detections=detections
+    ))

pixeltable/iterators/base.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
-from typing import Dict, Any, Tuple, List
-from abc import abstractmethod, ABC
+
+from abc import ABC, abstractmethod
+from typing import Any
 
 from pixeltable.type_system import ColumnType
 

pixeltable/iterators/document.py

@@ -152,7 +152,7 @@ class DocumentSplitter(ComponentIterator):
             assert self._doc_handle.pdf_doc is not None
             self._sections = self._pdf_sections()
         else:
-            assert False, f'unknown document format: {self._doc_handle.format}'
+            assert False, f'Unsupported document format: {self._doc_handle.format}'
 
         if Separator.SENTENCE in self._separators:
             self._sections = self._sentence_sections(self._sections)