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

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)

pixeltable/iterators/video.py

@@ -1,13 +1,15 @@
 import logging
 import math
+from fractions import Fraction
 from pathlib import Path
 from typing import Any, Optional, Sequence
 
-import cv2
+import av  # type: ignore[import-untyped]
+import pandas as pd
 import PIL.Image
 
-from pixeltable.exceptions import Error
-from pixeltable.type_system import ColumnType, FloatType, ImageType, IntType, VideoType
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
 
 from .base import ComponentIterator
 
@@ -30,108 +32,163 @@ class FrameIterator(ComponentIterator):
                 `num_frames` is greater than the number of frames in the video, all frames will be extracted.
     """
 
+    # Input parameters
     video_path: Path
-    video_reader: cv2.VideoCapture
     fps: Optional[float]
     num_frames: Optional[int]
-    frames_to_extract: Sequence[int]
-    frames_set: set[int]
-    next_frame_idx: int
+
+    # Video info
+    container: av.container.input.InputContainer
+    video_framerate: Fraction
+    video_time_base: Fraction
+    video_frame_count: int
+    video_start_time: int
+
+    # List of frame indices to be extracted, or None to extract all frames
+    frames_to_extract: Optional[list[int]]
+
+    # Next frame to extract, as an iterator `pos` index. If `frames_to_extract` is None, this is the same as the
+    # frame index in the video. Otherwise, the corresponding video index is `frames_to_extract[next_pos]`.
+    next_pos: int
 
     def __init__(self, video: str, *, fps: Optional[float] = None, num_frames: Optional[int] = None):
         if fps is not None and num_frames is not None:
-            raise Error('At most one of `fps` or `num_frames` may be specified')
+            raise excs.Error('At most one of `fps` or `num_frames` may be specified')
 
         video_path = Path(video)
         assert video_path.exists() and video_path.is_file()
         self.video_path = video_path
-        self.video_reader = cv2.VideoCapture(str(video_path))
+        self.container = av.open(str(video_path))
         self.fps = fps
         self.num_frames = num_frames
-        if not self.video_reader.isOpened():
-            raise Error(f'Failed to open video: {video}')
 
-        video_fps = int(self.video_reader.get(cv2.CAP_PROP_FPS))
-        if fps is not None and fps > video_fps:
-            raise Error(f'Video {video}: requested fps ({fps}) exceeds that of the video ({video_fps})')
-        num_video_frames = int(self.video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
-        if num_video_frames == 0:
-            raise Error(f'Video {video}: failed to get number of frames')
+        self.video_framerate = self.container.streams.video[0].average_rate
+        self.video_time_base = self.container.streams.video[0].time_base
+        self.video_start_time = self.container.streams.video[0].start_time or 0
+
+        # Determine the number of frames in the video
+        self.video_frame_count = self.container.streams.video[0].frames
+        if self.video_frame_count == 0:
+            # The video codec does not provide a frame count in the standard `frames` field. Try some other methods.
+            metadata: dict = self.container.streams.video[0].metadata
+            if 'NUMBER_OF_FRAMES' in metadata:
+                self.video_frame_count = int(metadata['NUMBER_OF_FRAMES'])
+            elif 'DURATION' in metadata:
+                # As a last resort, calculate the frame count from the stream duration.
+                duration = metadata['DURATION']
+                assert isinstance(duration, str)
+                seconds = pd.to_timedelta(duration).total_seconds()
+                # Usually the duration and framerate are precise enough for this calculation to be accurate, but if
+                # we encounter a case where it's off by one due to a rounding error, that's ok; we only use this
+                # to determine the positions of the sampled frames when `fps` or `num_frames` is specified.
+                self.video_frame_count = round(seconds * self.video_framerate)
+            else:
+                raise excs.Error(f'Video {video}: failed to get number of frames')
 
         if num_frames is not None:
             # specific number of frames
-            if num_frames > num_video_frames:
+            if num_frames > self.video_frame_count:
                 # Extract all frames
-                self.frames_to_extract = range(num_video_frames)
+                self.frames_to_extract = None
             else:
-                spacing = float(num_video_frames) / float(num_frames)
+                spacing = float(self.video_frame_count) / float(num_frames)
                 self.frames_to_extract = list(round(i * spacing) for i in range(num_frames))
                 assert len(self.frames_to_extract) == num_frames
         else:
             if fps is None or fps == 0.0:
                 # Extract all frames
-                self.frames_to_extract = range(num_video_frames)
+                self.frames_to_extract = None
+            elif fps > float(self.video_framerate):
+                raise excs.Error(
+                    f'Video {video}: requested fps ({fps}) exceeds that of the video ({float(self.video_framerate)})'
+                )
             else:
                 # Extract frames at the implied frequency
-                freq = fps / video_fps
-                n = math.ceil(num_video_frames * freq)  # number of frames to extract
+                freq = fps / float(self.video_framerate)
+                n = math.ceil(self.video_frame_count * freq)  # number of frames to extract
                 self.frames_to_extract = list(round(i / freq) for i in range(n))
 
-        # We need the list of frames as both a list (for set_pos) and a set (for fast lookups when
-        # there are lots of frames)
-        self.frames_set = set(self.frames_to_extract)
         _logger.debug(f'FrameIterator: path={self.video_path} fps={self.fps} num_frames={self.num_frames}')
-        self.next_frame_idx = 0
+        self.next_pos = 0
 
     @classmethod
-    def input_schema(cls) -> dict[str, ColumnType]:
+    def input_schema(cls) -> dict[str, ts.ColumnType]:
         return {
-            'video': VideoType(nullable=False),
-            'fps': FloatType(nullable=True),
-            'num_frames': IntType(nullable=True),
+            'video': ts.VideoType(nullable=False),
+            'fps': ts.FloatType(nullable=True),
+            'num_frames': ts.IntType(nullable=True),
         }
 
     @classmethod
-    def output_schema(cls, *args: Any, **kwargs: Any) -> tuple[dict[str, ColumnType], list[str]]:
+    def output_schema(cls, *args: Any, **kwargs: Any) -> tuple[dict[str, ts.ColumnType], list[str]]:
         return {
-            'frame_idx': IntType(),
-            'pos_msec': FloatType(),
-            'pos_frame': FloatType(),
-            'frame': ImageType(),
+            'frame_idx': ts.IntType(),
+            'pos_msec': ts.FloatType(),
+            'pos_frame': ts.IntType(),
+            'frame': ts.ImageType(),
         }, ['frame']
 
     def __next__(self) -> dict[str, Any]:
-        # jumping to the target frame here with video_reader.set() is far slower than just
-        # skipping the unwanted frames
+        # Determine the frame index in the video corresponding to the iterator index `next_pos`;
+        # the frame at this index is the one we want to extract next
+        if self.frames_to_extract is None:
+            next_video_idx = self.next_pos  # we're extracting all frames
+        elif self.next_pos >= len(self.frames_to_extract):
+            raise StopIteration
+        else:
+            next_video_idx = self.frames_to_extract[self.next_pos]
+
+        # We are searching for the frame at the index implied by `next_pos`. Step through the video until we
+        # find it. There are two reasons why it might not be the immediate next frame in the video:
+        # (1) `fps` or `num_frames` was specified as an iterator argument; or
+        # (2) we just did a seek, and the desired frame is not a keyframe.
+        # TODO: In case (1) it will usually be fastest to step through the frames until we find the one we're
+        #     looking for. But in some cases it may be faster to do a seek; for example, when `fps` is very
+        #     low and there are multiple keyframes in between each frame we want to extract (imagine extracting
+        #     10 frames from an hourlong video).
         while True:
-            pos_msec = self.video_reader.get(cv2.CAP_PROP_POS_MSEC)
-            pos_frame = self.video_reader.get(cv2.CAP_PROP_POS_FRAMES)
-            status, img = self.video_reader.read()
-            if not status:
-                _logger.debug(f'releasing video reader for {self.video_path}')
-                self.video_reader.release()
-                self.video_reader = None
+            try:
+                frame = next(self.container.decode(video=0))
+            except EOFError:
                 raise StopIteration
-            if pos_frame in self.frames_set:
-                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-                result = {
-                    'frame_idx': self.next_frame_idx,
-                    'pos_msec': pos_msec,
-                    'pos_frame': pos_frame,
-                    'frame': PIL.Image.fromarray(img),
-                }
-                self.next_frame_idx += 1
-                return result
+            # Compute the index of the current frame in the video based on the presentation timestamp (pts);
+            # this ensures we have a canonical understanding of frame index, regardless of how we got here
+            # (seek or iteration)
+            pts = frame.pts - self.video_start_time
+            video_idx = round(pts * self.video_time_base * self.video_framerate)
+            assert isinstance(video_idx, int)
+            if video_idx < next_video_idx:
+                # We haven't reached the desired frame yet
+                continue
+
+            # Sanity check that we're at the right frame.
+            if video_idx != next_video_idx:
+                raise excs.Error(f'Frame {next_video_idx} is missing from the video (video file is corrupt)')
+            img = frame.to_image()
+            assert isinstance(img, PIL.Image.Image)
+            pos_msec = float(pts * self.video_time_base * 1000)
+            result = {
+                'frame_idx': self.next_pos,
+                'pos_msec': pos_msec,
+                'pos_frame': video_idx,
+                'frame': img,
+            }
+            self.next_pos += 1
+            return result
 
     def close(self) -> None:
-        if self.video_reader is not None:
-            self.video_reader.release()
-            self.video_reader = None
+        self.container.close()
 
     def set_pos(self, pos: int) -> None:
         """Seek to frame idx"""
-        if pos == self.next_frame_idx:
-            return
-        _logger.debug(f'seeking to frame {pos}')
-        self.video_reader.set(cv2.CAP_PROP_POS_FRAMES, self.frames_to_extract[pos])
-        self.next_frame_idx = pos
+        if pos == self.next_pos:
+            return  # already there
+
+        video_idx = pos if self.frames_to_extract is None else self.frames_to_extract[pos]
+        _logger.debug(f'seeking to frame number {video_idx} (at iterator index {pos})')
+        # compute the frame position in time_base units
+        seek_pos = int(video_idx / self.video_framerate / self.video_time_base + self.video_start_time)
+        # This will seek to the nearest keyframe before the desired frame. If the frame being sought is not a keyframe,
+        # then the iterator will step forward to the desired frame on the subsequent call to next().
+        self.container.seek(seek_pos, backward=True, stream=self.container.streams.video[0])
+        self.next_pos = pos

pixeltable/metadata/__init__.py

@@ -10,7 +10,7 @@ import sqlalchemy.orm as orm
 from .schema import SystemInfo, SystemInfoMd
 
 # current version of the metadata; this is incremented whenever the metadata schema changes
-VERSION = 21
+VERSION = 22
 
 
 def create_system_info(engine: sql.engine.Engine) -> None:

pixeltable/metadata/converters/convert_21.py

@@ -0,0 +1,34 @@
+from typing import Any, Optional
+import sqlalchemy as sql
+
+from pixeltable.metadata import register_converter
+from pixeltable.metadata.converters.util import convert_table_schema_version_md, convert_table_md
+
+
+@register_converter(version=21)
+def _(engine: sql.engine.Engine) -> None:
+    convert_table_schema_version_md(
+        engine,
+        table_schema_version_md_updater=__update_table_schema_version,
+        schema_column_updater=__update_schema_column
+    )
+    convert_table_md(
+        engine,
+        substitution_fn=__substitute_md
+    )
+
+
+def __update_table_schema_version(table_schema_version_md: dict) -> None:
+    table_schema_version_md['media_validation'] = 'on_write'  # MediaValidation.ON_WRITE
+
+
+def __update_schema_column(schema_column: dict) -> None:
+    schema_column['media_validation'] = None
+
+
+def __substitute_md(k: Optional[str], v: Any) -> Optional[tuple[Optional[str], Any]]:
+    if isinstance(v, dict) and '_classname' in v and v['_classname'] == 'ColumnRef':
+        if 'perform_validation' not in v:
+            v['perform_validation'] = False
+        return k, v
+    return None