pixeltable 0.2.18__py3-none-any.whl → 0.2.19__py3-none-any.whl

pixeltable/__init__.py

@@ -4,7 +4,7 @@ from .exceptions import Error
 from .exprs import RELATIVE_PATH_ROOT
 from .func import Function, udf, Aggregator, uda, expr_udf
 from .globals import init, create_table, create_view, get_table, move, drop_table, list_tables, create_dir, drop_dir, \
-    list_dirs, list_functions, configure_logging
+    list_dirs, list_functions, configure_logging, array
 from .type_system import (
     ColumnType,
     StringType,

pixeltable/__version__.py

@@ -1,3 +1,3 @@
 # These version placeholders will be replaced during build.
-__version__ = "0.2.18"
-__version_tuple__ = (0, 2, 18)
+__version__ = "0.2.19"
+__version_tuple__ = (0, 2, 19)

pixeltable/catalog/table.py

@@ -2,7 +2,6 @@ from __future__ import annotations
 
 import abc
 import builtins
-import itertools
 import json
 import logging
 from pathlib import Path

pixeltable/catalog/table_version.py

@@ -147,7 +147,7 @@ class TableVersion:
             module = importlib.import_module(module_name)
             self.iterator_cls = getattr(module, class_name)
             self.iterator_args = exprs.InlineDict.from_dict(tbl_md.view_md.iterator_args)
-            output_schema, _ = self.iterator_cls.output_schema(**self.iterator_args.to_dict())
+            output_schema, _ = self.iterator_cls.output_schema(**self.iterator_args.to_kwargs())
             self.num_iterator_cols = len(output_schema)
             assert tbl_md.view_md.iterator_args is not None
 

pixeltable/catalog/view.py

@@ -94,7 +94,7 @@ class View(Table):
                 ]
                 sig = func.Signature(InvalidType(), params)
                 from pixeltable.exprs import FunctionCall
-                FunctionCall.normalize_args(sig, bound_args)
+                FunctionCall.normalize_args(iterator_cls.__name__, sig, bound_args)
             except TypeError as e:
                 raise Error(f'Cannot instantiate iterator with given arguments: {e}')
 

pixeltable/dataframe.py

@@ -501,7 +501,7 @@ class DataFrame:
             elif isinstance(raw_expr, dict):
                 select_list.append((exprs.InlineDict(raw_expr), name))
             elif isinstance(raw_expr, list):
-                select_list.append((exprs.InlineArray(raw_expr), name))
+                select_list.append((exprs.InlineList(raw_expr), name))
             else:
                 select_list.append((exprs.Literal(raw_expr), name))
             expr = select_list[-1][0]

pixeltable/env.py

@@ -8,6 +8,7 @@ import importlib.util
 import inspect
 import logging
 import os
+import subprocess
 import sys
 import threading
 import uuid
@@ -486,6 +487,8 @@ class Env:
         """Check for and start runtime services"""
         self._start_web_server()
         self.__register_packages()
+        if self.is_installed_package('spacy'):
+            self.__init_spacy()
 
     def __register_packages(self) -> None:
         """Declare optional packages that are utilized by some parts of the code."""
@@ -500,7 +503,7 @@ class Env:
         self.__register_package('openpyxl')
         self.__register_package('pyarrow')
         self.__register_package('sentence_transformers', library_name='sentence-transformers')
-        self.__register_package('spacy')  # TODO: deal with en-core-web-sm
+        self.__register_package('spacy')
         self.__register_package('tiktoken')
         self.__register_package('together')
         self.__register_package('toml')
@@ -511,10 +514,6 @@ class Env:
         self.__register_package('whisperx')
         self.__register_package('yolox', library_name='git+https://github.com/Megvii-BaseDetection/YOLOX@ac58e0a')
 
-        if self.is_installed_package('spacy'):
-            import spacy
-            self._spacy_nlp = spacy.load('en_core_web_sm')
-
     def __register_package(self, package_name: str, library_name: Optional[str] = None) -> None:
         self.__optional_packages[package_name] = PackageInfo(
             is_installed=importlib.util.find_spec(package_name) is not None,
@@ -556,6 +555,35 @@ class Env:
                 f'To fix this, run: `pip install -U {package_info.library_name}`'
             )
 
+    def __init_spacy(self) -> None:
+        """
+        spaCy relies on a pip-installed model to operate. In order to avoid requiring the model as a separate
+        dependency, we install it programmatically here. This should cause no problems, since the model packages
+        have no sub-dependencies (in fact, this is how spaCy normally manages its model resources).
+        """
+        import spacy
+        from spacy.cli.download import get_model_filename
+        spacy_model = 'en_core_web_sm'
+        spacy_model_version = '3.7.1'
+        filename = get_model_filename(spacy_model, spacy_model_version, sdist=False)
+        url = f'{spacy.about.__download_url__}/{filename}'
+        # Try to `pip install` the model. We set check=False; if the pip command fails, it's not necessarily
+        # a problem, because the model have been installed on a previous attempt.
+        self._logger.info(f'Ensuring spaCy model is installed: {filename}')
+        ret = subprocess.run([sys.executable, '-m', 'pip', 'install', '-qU', url], check=False)
+        if ret.returncode != 0:
+            self._logger.warn(f'pip install failed for spaCy model: {filename}')
+        try:
+            self._logger.info(f'Loading spaCy model: {spacy_model}')
+            self._spacy_nlp = spacy.load(spacy_model)
+        except Exception as exc:
+            self._logger.warn(f'Failed to load spaCy model: {spacy_model}', exc_info=exc)
+            warnings.warn(
+                f"Failed to load spaCy model '{spacy_model}'. spaCy features will not be available.",
+                excs.PixeltableWarning
+            )
+            self.__optional_packages['spacy'].is_installed = False
+
     def num_tmp_files(self) -> int:
         return len(glob.glob(f'{self._tmp_dir}/*'))
 
@@ -594,6 +622,7 @@ class Env:
 
     @property
     def spacy_nlp(self) -> spacy.Language:
+        Env.get().require_package('spacy')
         assert self._spacy_nlp is not None
         return self._spacy_nlp
 

pixeltable/exceptions.py

@@ -14,4 +14,8 @@ class ExprEvalError(Exception):
     exc: Exception
     exc_tb: TracebackType
     input_vals: List[Any]
-    row_num: int
+    row_num: int
+
+
+class PixeltableWarning(Warning):
+    pass

pixeltable/exec/component_iteration_node.py

@@ -24,7 +24,7 @@ class ComponentIterationNode(ExecNode):
         assert isinstance(self.iterator_args, exprs.InlineDict)
         self.iterator_args_ctx = self.row_builder.create_eval_ctx([self.iterator_args])
         self.iterator_output_schema, self.unstored_column_names = \
-            self.view.iterator_cls.output_schema(**self.iterator_args.to_dict())
+            self.view.iterator_cls.output_schema(**self.iterator_args.to_kwargs())
         self.iterator_output_fields = list(self.iterator_output_schema.keys())
         self.iterator_output_cols = \
             {field_name: self.view.cols_by_name[field_name] for field_name in self.iterator_output_fields}

pixeltable/exprs/__init__.py

@@ -9,8 +9,7 @@ from .expr import Expr
 from .expr_set import ExprSet
 from .function_call import FunctionCall
 from .in_predicate import InPredicate
-from .inline_array import InlineArray
-from .inline_dict import InlineDict
+from .inline_expr import InlineArray, InlineDict, InlineList
 from .is_null import IsNull
 from .json_mapper import JsonMapper
 from .json_path import RELATIVE_PATH_ROOT, JsonPath

pixeltable/exprs/expr.py

@@ -356,15 +356,14 @@ class Expr(abc.ABC):
         """
         if isinstance(o, Expr):
             return o
-        # Try to create a literal. We need to check for InlineArray/InlineDict
-        # first, to prevent arrays from inappropriately being interpreted as JsonType
+        # Try to create a literal. We need to check for InlineList/InlineDict
+        # first, to prevent them from inappropriately being interpreted as JsonType
         # literals.
-        # TODO: general cleanup of InlineArray/InlineDict
         if isinstance(o, list):
-            from .inline_array import InlineArray
-            return InlineArray(tuple(o))
+            from .inline_expr import InlineList
+            return InlineList(o)
         if isinstance(o, dict):
-            from .inline_dict import InlineDict
+            from .inline_expr import InlineDict
             return InlineDict(o)
         obj_type = ts.ColumnType.infer_literal_type(o)
         if obj_type is not None:

pixeltable/exprs/function_call.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import inspect
 import json
 import sys
-from typing import Optional, Any
+from typing import Any, Optional
 
 import sqlalchemy as sql
 
@@ -11,10 +11,10 @@ import pixeltable.catalog as catalog
 import pixeltable.exceptions as excs
 import pixeltable.func as func
 import pixeltable.type_system as ts
+
 from .data_row import DataRow
 from .expr import Expr
-from .inline_array import InlineArray
-from .inline_dict import InlineDict
+from .inline_expr import InlineDict, InlineList
 from .row_builder import RowBuilder
 from .rowid_ref import RowidRef
 from .sql_element_cache import SqlElementCache
@@ -53,7 +53,7 @@ class FunctionCall(Expr):
         super().__init__(fn.call_return_type(bound_args))
         self.fn = fn
         self.is_method_call = is_method_call
-        self.normalize_args(signature, bound_args)
+        self.normalize_args(fn.name, signature, bound_args)
 
         self.agg_init_args = {}
         if self.is_agg_fn_call:
@@ -143,7 +143,7 @@ class FunctionCall(Expr):
         return super().default_column_name()
 
     @classmethod
-    def normalize_args(cls, signature: func.Signature, bound_args: dict[str, Any]) -> None:
+    def normalize_args(cls, fn_name: str, signature: func.Signature, bound_args: dict[str, Any]) -> None:
         """Converts all args to Exprs and checks that they are compatible with signature.
 
         Updates bound_args in place, where necessary.
@@ -163,9 +163,7 @@ class FunctionCall(Expr):
 
             if isinstance(arg, list) or isinstance(arg, tuple):
                 try:
-                    # If the column type is JsonType, force the literal to be JSON
-                    is_json = is_var_param or (param.col_type is not None and param.col_type.is_json_type())
-                    arg = InlineArray(arg, force_json=is_json)
+                    arg = InlineList(arg)
                     bound_args[param_name] = arg
                     continue
                 except excs.Error:
@@ -177,7 +175,7 @@ class FunctionCall(Expr):
                 try:
                     _ = json.dumps(arg)
                 except TypeError:
-                    raise excs.Error(f'Argument for parameter {param_name!r} is not json-serializable: {arg}')
+                    raise excs.Error(f'Argument for parameter {param_name!r} is not json-serializable: {arg} (of type {type(arg)})')
                 if arg is not None:
                     try:
                         param_type = param.col_type
@@ -215,7 +213,7 @@ class FunctionCall(Expr):
                     or (arg.col_type.is_json_type() and param.col_type.is_scalar_type())
                 ):
                     raise excs.Error(
-                        f'Parameter {param_name}: argument type {arg.col_type} does not match parameter type '
+                        f'Parameter {param_name} (in function {fn_name}): argument type {arg.col_type} does not match parameter type '
                         f'{param.col_type}')
 
     def _equals(self, other: FunctionCall) -> bool:

pixeltable/exprs/inline_expr.py

@@ -0,0 +1,200 @@
+from __future__ import annotations
+
+import copy
+from typing import Any, Iterable, Optional
+
+import numpy as np
+import sqlalchemy as sql
+
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+
+from .data_row import DataRow
+from .expr import Expr
+from .literal import Literal
+from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
+
+
+class InlineArray(Expr):
+    """
+    Array 'literal' which can use Exprs as values.
+    """
+
+    def __init__(self, elements: Iterable):
+        exprs = []
+        for el in elements:
+            if isinstance(el, Expr):
+                exprs.append(el)
+            elif isinstance(el, list) or isinstance(el, tuple):
+                exprs.append(InlineArray(el))
+            else:
+                exprs.append(Literal(el))
+
+        inferred_element_type: Optional[ts.ColumnType] = ts.InvalidType()
+        for i, expr in enumerate(exprs):
+            supertype = inferred_element_type.supertype(expr.col_type)
+            if supertype is None:
+                raise excs.Error(
+                    f'Could not infer element type of array: element of type `{expr.col_type}` at index {i} '
+                    f'is not compatible with type `{inferred_element_type}` of preceding elements'
+                )
+            inferred_element_type = supertype
+
+        if inferred_element_type.is_scalar_type():
+            col_type = ts.ArrayType((len(exprs),), inferred_element_type)
+        elif inferred_element_type.is_array_type():
+            assert isinstance(inferred_element_type, ts.ArrayType)
+            col_type = ts.ArrayType(
+                (len(exprs), *inferred_element_type.shape),
+                ts.ColumnType.make_type(inferred_element_type.dtype)
+            )
+        else:
+            raise excs.Error(f'Element type is not a valid dtype for an array: {inferred_element_type}')
+
+        super().__init__(col_type)
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        elem_strs = [str(expr) for expr in self.components]
+        return f'[{", ".join(elem_strs)}]'
+
+    def _equals(self, _: InlineArray) -> bool:
+        return True  # Always true if components match
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        data_row[self.slot_idx] = np.array([data_row[el.slot_idx] for el in self.components])
+
+    def _as_dict(self) -> dict:
+        return super()._as_dict()
+
+    @classmethod
+    def _from_dict(cls, _: dict, components: list[Expr]) -> Expr:
+        try:
+            return cls(components)
+        except excs.Error:
+            # For legacy compatibility reasons, we need to try constructing as an `InlineList`.
+            # This is because in schema versions <= 19, `InlineArray` was serialized incorrectly, and
+            # there is no way to determine the correct expression type until the subexpressions are
+            # loaded and their types are known.
+            return InlineList(components)
+
+
+class InlineList(Expr):
+    """
+    List 'literal' which can use Exprs as values.
+    """
+
+    def __init__(self, elements: Iterable):
+        exprs = []
+        for el in elements:
+            if isinstance(el, Expr):
+                exprs.append(el)
+            elif isinstance(el, list) or isinstance(el, tuple):
+                exprs.append(InlineList(el))
+            elif isinstance(el, dict):
+                exprs.append(InlineDict(el))
+            else:
+                exprs.append(Literal(el))
+
+        super().__init__(ts.JsonType())
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        elem_strs = [str(expr) for expr in self.components]
+        return f'[{", ".join(elem_strs)}]'
+
+    def _equals(self, _: InlineList) -> bool:
+        return True  # Always true if components match
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, _: RowBuilder) -> None:
+        data_row[self.slot_idx] = [data_row[el.slot_idx] for el in self.components]
+
+    def _as_dict(self) -> dict:
+        return super()._as_dict()
+
+    @classmethod
+    def _from_dict(cls, _: dict, components: list[Expr]) -> Expr:
+        return cls(components)
+
+
+class InlineDict(Expr):
+    """
+    Dictionary 'literal' which can use Exprs as values.
+    """
+
+    keys: list[str]
+
+    def __init__(self, d: dict[str, Any]):
+        self.keys = []
+        exprs: list[Expr] = []
+        for key, val in d.items():
+            if not isinstance(key, str):
+                raise excs.Error(f'Dictionary requires string keys; {key} has type {type(key)}')
+            self.keys.append(key)
+            if isinstance(val, Expr):
+                exprs.append(val)
+            elif isinstance(val, dict):
+                exprs.append(InlineDict(val))
+            elif isinstance(val, list) or isinstance(val, tuple):
+                exprs.append(InlineList(val))
+            else:
+                exprs.append(Literal(val))
+
+        super().__init__(ts.JsonType())
+        self.components.extend(exprs)
+        self.id = self._create_id()
+
+    def __str__(self) -> str:
+        item_strs = list(f"'{key}': {str(expr)}" for key, expr in zip(self.keys, self.components))
+        return '{' + ', '.join(item_strs) + '}'
+
+    def _equals(self, other: InlineDict) -> bool:
+        # The dict values are just the components, which have already been checked
+        return self.keys == other.keys
+
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        return super()._id_attrs() + [('keys', self.keys)]
+
+    def sql_expr(self, _: SqlElementCache) -> Optional[sql.ColumnElement]:
+        return None
+
+    def eval(self, data_row: DataRow, _: RowBuilder) -> None:
+        assert len(self.keys) == len(self.components)
+        data_row[self.slot_idx] = {
+            key: data_row[expr.slot_idx]
+            for key, expr in zip(self.keys, self.components)
+        }
+
+    def to_kwargs(self) -> dict[str, Any]:
+        """Deconstructs this expression into a dictionary by recursively unwrapping all Literals,
+        InlineDicts, and InlineLists."""
+        return InlineDict._to_kwarg_element(self)
+
+    @classmethod
+    def _to_kwarg_element(cls, expr: Expr) -> Any:
+        if isinstance(expr, Literal):
+            return expr.val
+        if isinstance(expr, InlineDict):
+            return {key: cls._to_kwarg_element(val) for key, val in zip(expr.keys, expr.components)}
+        if isinstance(expr, InlineList):
+            return [cls._to_kwarg_element(el) for el in expr.components]
+        return expr
+
+    def _as_dict(self) -> dict[str, Any]:
+        return {'keys': self.keys, **super()._as_dict()}
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: list[Expr]) -> Expr:
+        assert 'keys' in d
+        assert len(d['keys']) == len(components)
+        arg = dict(zip(d['keys'], components))
+        return InlineDict(arg)

pixeltable/ext/functions/whisperx.py

@@ -19,6 +19,8 @@ def transcribe(
     equivalent to the WhisperX `transcribe` function, as described in the
     [WhisperX library documentation](https://github.com/m-bain/whisperX).
 
+    WhisperX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
+
     __Requirements:__
 
     - `pip install whisperx`

pixeltable/ext/functions/yolox.py

@@ -26,8 +26,7 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
     Computes YOLOX object detections for the specified image. `model_id` should reference one of the models
     defined in the [YOLOX documentation](https://github.com/Megvii-BaseDetection/YOLOX).
 
-    YOLOX support is part of the `pixeltable.ext` package: long-term support is not guaranteed, and it is not
-    intended for use in production applications.
+    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
 
     __Requirements__:
 
@@ -79,6 +78,8 @@ def yolo_to_coco(detections: dict) -> list:
     """
     Converts the output of a YOLOX object detection model to COCO format.
 
+    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
+
     Args:
         detections: The output of a YOLOX object detection model, as returned by `yolox`.
 
@@ -89,7 +90,8 @@ def yolo_to_coco(detections: dict) -> list:
         Add a computed column that converts the output `tbl.detections` to COCO format, where `tbl.image`
         is the image for which detections were computed:
 
-        >>> tbl['detections_coco'] = yolo_to_coco(tbl.detections)
+        >>> tbl['detections'] = yolox(tbl.image, model_id='yolox_m', threshold=0.8)
+        ... tbl['detections_coco'] = yolo_to_coco(tbl.detections)
     """
     bboxes, labels = detections['bboxes'], detections['labels']
     num_annotations = len(detections['bboxes'])

pixeltable/functions/huggingface.py

@@ -185,7 +185,7 @@ def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndar
 
     Examples:
         Add a computed column that applies the model `openai/clip-vit-base-patch32` to an existing
-        Pixeltable column `tbl.image` of the table `tbl`:
+        Pixeltable column `image` of the table `tbl`:
 
         >>> tbl['result'] = clip_image(tbl.image, model_id='openai/clip-vit-base-patch32')
     """
@@ -228,24 +228,24 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
 
     Args:
         image: The image to embed.
-        model_id: The pretrained model to use for the embedding.
+        model_id: The pretrained model to use for object detection.
 
     Returns:
         A dictionary containing the output of the object detection model, in the following format:
 
-    ```python
-    {
-        'scores': [0.99, 0.999],  # list of confidence scores for each detected object
-        'labels': [25, 25],  # list of COCO class labels for each detected object
-        'label_text': ['giraffe', 'giraffe'],  # corresponding text names of class labels
-        'boxes': [[51.942, 356.174, 181.481, 413.975], [383.225, 58.66, 605.64, 361.346]]
-            # list of bounding boxes for each detected object, as [x1, y1, x2, y2]
-    }
-    ```
+            ```python
+            {
+                'scores': [0.99, 0.999],  # list of confidence scores for each detected object
+                'labels': [25, 25],  # list of COCO class labels for each detected object
+                'label_text': ['giraffe', 'giraffe'],  # corresponding text names of class labels
+                'boxes': [[51.942, 356.174, 181.481, 413.975], [383.225, 58.66, 605.64, 361.346]]
+                    # list of bounding boxes for each detected object, as [x1, y1, x2, y2]
+            }
+            ```
 
     Examples:
         Add a computed column that applies the model `facebook/detr-resnet-50` to an existing
-        Pixeltable column `tbl.image` of the table `tbl`:
+        Pixeltable column `image` of the table `tbl`:
 
         >>> tbl['detections'] = detr_for_object_detection(
         ...     tbl.image,
@@ -282,6 +282,83 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
     ]
 
 
+@pxt.udf(batch_size=4)
+def vit_for_image_classification(
+    image: Batch[PIL.Image.Image],
+    *,
+    model_id: str,
+    top_k: int = 5
+) -> Batch[list[dict[str, Any]]]:
+    """
+    Computes image classifications for the specified image using a Vision Transformer (ViT) model.
+    `model_id` should be a reference to a pretrained [ViT Model](https://huggingface.co/docs/transformers/en/model_doc/vit).
+
+    __Note:__ Be sure the model is a ViT model that is trained for image classification (that is, a model designed for
+    use with the
+    [ViTForImageClassification](https://huggingface.co/docs/transformers/en/model_doc/vit#transformers.ViTForImageClassification)
+    class), such as `google/vit-base-patch16-224`. General feature-extraction models such as
+    `google/vit-base-patch16-224-in21k` will not produce the desired results.
+
+    __Requirements:__
+
+    - `pip install transformers`
+
+    Args:
+        image: The image to classify.
+        model_id: The pretrained model to use for the classification.
+        top_k: The number of classes to return.
+
+    Returns:
+        A list of the `top_k` highest-scoring classes for each image. Each element in the list is a dictionary
+            in the following format:
+
+            ```python
+            {
+                'p': 0.230,  # class probability
+                'class': 935,  # class ID
+                'label': 'mashed potato',  # class label
+            }
+            ```
+
+    Examples:
+        Add a computed column that applies the model `google/vit-base-patch16-224` to an existing
+        Pixeltable column `image` of the table `tbl`:
+
+        >>> tbl['image_class'] = vit_for_image_classification(
+        ...     tbl.image,
+        ...     model_id='google/vit-base-patch16-224'
+        ... )
+    """
+    env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
+    from transformers import ViTImageProcessor, ViTForImageClassification
+
+    model: ViTForImageClassification = _lookup_model(model_id, ViTForImageClassification.from_pretrained, device=device)
+    processor = _lookup_processor(model_id, ViTImageProcessor.from_pretrained)
+    normalized_images = [normalize_image_mode(img) for img in image]
+
+    with torch.no_grad():
+        inputs = processor(images=normalized_images, return_tensors='pt')
+        outputs = model(**inputs.to(device))
+        logits = outputs.logits
+
+    probs = torch.softmax(logits, dim=-1)
+    top_k_probs, top_k_indices = torch.topk(probs, top_k, dim=-1)
+
+    return [
+        [
+            {
+                'p': top_k_probs[n, k].item(),
+                'class': top_k_indices[n, k].item(),
+                'label': model.config.id2label[top_k_indices[n, k].item()],
+            }
+            for k in range(top_k_probs.shape[1])
+        ]
+        for n in range(top_k_probs.shape[0])
+    ]
+
+
 @pxt.udf
 def detr_to_coco(image: PIL.Image.Image, detr_info: dict[str, Any]) -> dict[str, Any]:
     """

pixeltable/functions/image.py

@@ -92,7 +92,7 @@ def _(self: Expr, mode: str) -> ts.ColumnType:
 
 
 # Image.crop()
-@func.udf(substitute_fn=PIL.Image.Image.crop, param_types=[ts.ImageType(), ts.ArrayType((4,), dtype=ts.IntType())], is_method=True)
+@func.udf(substitute_fn=PIL.Image.Image.crop, is_method=True)
 def crop(self: PIL.Image.Image, box: tuple[int, int, int, int]) -> PIL.Image.Image:
     """
     Return a rectangular region from the image. The box is a 4-tuple defining the left, upper, right, and lower pixel
@@ -151,7 +151,7 @@ def _(self: Expr) -> ts.ColumnType:
 
 
 # Image.resize()
-@func.udf(param_types=[ts.ImageType(), ts.ArrayType((2,), dtype=ts.IntType())], is_method=True)
+@func.udf(is_method=True)
 def resize(self: PIL.Image.Image, size: tuple[int, int]) -> PIL.Image.Image:
     """
     Return a resized copy of the image. The size parameter is a tuple containing the width and height of the new image.
@@ -366,7 +366,7 @@ def quantize(
 
 
 @func.udf(substitute_fn=PIL.Image.Image.reduce, is_method=True)
-def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int]] = None) -> PIL.Image.Image:
+def reduce(self: PIL.Image.Image, factor: int, box: Optional[tuple[int, int, int, int]] = None) -> PIL.Image.Image:
     """
     Reduce the image by the given factor.