pixeltable 0.3.9__py3-none-any.whl → 0.3.11__py3-none-any.whl

pixeltable/ext/functions/yolox.py

@@ -1,21 +1,15 @@
 import logging
-from pathlib import Path
-from typing import TYPE_CHECKING, Iterable, Iterator
-from urllib.request import urlretrieve
+from typing import TYPE_CHECKING
 
-import numpy as np
 import PIL.Image
 
 import pixeltable as pxt
-from pixeltable import env
 from pixeltable.func import Batch
 from pixeltable.functions.util import normalize_image_mode
 from pixeltable.utils.code import local_public_names
 
 if TYPE_CHECKING:
-    import torch
-    from yolox.exp import Exp  # type: ignore[import-untyped]
-    from yolox.models import YOLOX  # type: ignore[import-untyped]
+    from yolox.models import Yolox, YoloxProcessor  # type: ignore[import-untyped]
 
 _logger = logging.getLogger('pixeltable')
 
@@ -30,7 +24,7 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
 
     __Requirements__:
 
-    - `pip install git+https://github.com/Megvii-BaseDetection/YOLOX`
+    - `pip install pixeltable-yolox`
 
     Args:
         model_id: one of: `yolox_nano`, `yolox_tiny`, `yolox_s`, `yolox_m`, `yolox_l`, `yolox_x`
@@ -46,31 +40,14 @@ def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0
         >>> tbl.add_computed_column(detections=yolox(tbl.image, model_id='yolox_m', threshold=0.8))
     """
     import torch
-    from yolox.utils import postprocess  # type: ignore[import-untyped]
-
-    model, exp = _lookup_model(model_id, 'cpu')
-    image_tensors = list(_images_to_tensors(images, exp))
-    batch_tensor = torch.stack(image_tensors)
 
+    model = _lookup_model(model_id, 'cpu')
+    processor = _lookup_processor(model_id)
+    normalized_images = [normalize_image_mode(image) for image in images]
     with torch.no_grad():
-        output_tensor = model(batch_tensor)
-
-    outputs = postprocess(output_tensor, 80, threshold, exp.nmsthre, class_agnostic=False)
-
-    results: list[dict] = []
-    for image in images:
-        ratio = min(exp.test_size[0] / image.height, exp.test_size[1] / image.width)
-        if outputs[0] is None:
-            results.append({'bboxes': [], 'scores': [], 'labels': []})
-        else:
-            results.append(
-                {
-                    'bboxes': [(output[:4] / ratio).tolist() for output in outputs[0]],
-                    'scores': [output[4].item() * output[5].item() for output in outputs[0]],
-                    'labels': [int(output[6]) for output in outputs[0]],
-                }
-            )
-    return results
+        tensor = processor(normalized_images)
+        output = model(tensor)
+        return processor.postprocess(normalized_images, output, threshold=threshold)
 
 
 @pxt.udf
@@ -107,51 +84,31 @@ def yolo_to_coco(detections: dict) -> list:
     return result
 
 
-def _images_to_tensors(images: Iterable[PIL.Image.Image], exp: 'Exp') -> Iterator['torch.Tensor']:
-    import torch
-    from yolox.data import ValTransform  # type: ignore[import-untyped]
-
-    val_transform = ValTransform(legacy=False)
-    for image in images:
-        normalized_image = normalize_image_mode(image)
-        image_transform, _ = val_transform(np.array(normalized_image), None, exp.test_size)
-        yield torch.from_numpy(image_transform)
-
-
-def _lookup_model(model_id: str, device: str) -> tuple['YOLOX', 'Exp']:
-    import torch
-    from yolox.exp import get_exp
+def _lookup_model(model_id: str, device: str) -> 'Yolox':
+    from yolox.models import Yolox
 
     key = (model_id, device)
-    if key in _model_cache:
-        return _model_cache[key]
+    if key not in _model_cache:
+        _model_cache[key] = Yolox.from_pretrained(model_id, device=device)
 
-    weights_url = f'https://github.com/Megvii-BaseDetection/YOLOX/releases/download/0.1.1rc0/{model_id}.pth'
-    weights_file = Path(f'{env.Env.get().tmp_dir}/{model_id}.pth')
-    if not weights_file.exists():
-        _logger.info(f'Downloading weights for YOLOX model {model_id}: from {weights_url} -> {weights_file}')
-        urlretrieve(weights_url, weights_file)
+    return _model_cache[key]
 
-    exp = get_exp(exp_name=model_id)
-    model = exp.get_model().to(device)
 
-    model.eval()
-    model.head.training = False
-    model.training = False
+def _lookup_processor(model_id: str) -> 'YoloxProcessor':
+    from yolox.models import YoloxProcessor
 
-    # Load in the weights from training
-    weights = torch.load(weights_file, map_location=torch.device(device))
-    model.load_state_dict(weights['model'])
+    if model_id not in _processor_cache:
+        _processor_cache[model_id] = YoloxProcessor(model_id)
 
-    _model_cache[key] = (model, exp)
-    return model, exp
+    return _processor_cache[model_id]
 
 
-_model_cache: dict[tuple[str, str], tuple['YOLOX', 'Exp']] = {}
+_model_cache: dict[tuple[str, str], 'Yolox'] = {}
+_processor_cache: dict[str, 'YoloxProcessor'] = {}
 
 
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/func/aggregate_function.py

@@ -252,7 +252,7 @@ def uda(
 ) -> Callable[[type[Aggregator]], AggregateFunction]: ...
 
 
-def uda(*args, **kwargs):
+def uda(*args, **kwargs):  # type: ignore[no-untyped-def]
     """Decorator for user-defined aggregate functions.
 
     The decorated class must inherit from Aggregator and implement the following methods:

pixeltable/func/callable_function.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import inspect
 from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence
 from uuid import UUID
@@ -127,12 +126,10 @@ class CallableFunction(Function):
         """
         assert self.is_batched
         assert not self.is_polymorphic
+        assert not self.is_async
         # Unpack the constant parameters
         constant_kwargs, batched_kwargs = self.create_batch_kwargs(kwargs)
-        if inspect.iscoroutinefunction(self.py_fn):
-            return asyncio.run(self.py_fn(*args, **constant_kwargs, **batched_kwargs))
-        else:
-            return self.py_fn(*args, **constant_kwargs, **batched_kwargs)
+        return self.py_fn(*args, **constant_kwargs, **batched_kwargs)
 
     def create_batch_kwargs(self, kwargs: dict[str, Any]) -> tuple[dict[str, Any], dict[str, list[Any]]]:
         """Converts kwargs containing lists into constant and batched kwargs in the format expected by a batched udf."""

pixeltable/func/expr_template_function.py

@@ -1,6 +1,6 @@
 from typing import Any, Optional, Sequence
 
-from pixeltable import exceptions as excs, exprs
+from pixeltable import exceptions as excs, exprs, type_system as ts
 
 from .function import Function
 from .signature import Signature
@@ -76,9 +76,25 @@ class ExprTemplateFunction(Function):
                 arg_expr = arg
             arg_exprs[param_expr] = arg_expr
         result = result.substitute(arg_exprs)
-        assert not result._contains(exprs.Variable)
         return result
 
+    def call_return_type(self, bound_args: dict[str, 'exprs.Expr']) -> ts.ColumnType:
+        """
+        The call_return_type of an ExprTemplateFunction is derived from the template expression's col_type after
+        substitution (unlike for UDFs, whose call_return_type is derived from an explicitly specified
+        conditional_return_type).
+        """
+        assert not self.is_polymorphic
+        template = self.template
+        with_defaults = bound_args.copy()
+        with_defaults.update(
+            {param_name: default for param_name, default in template.defaults.items() if param_name not in bound_args}
+        )
+        substituted_expr = self.template.expr.copy().substitute(
+            {template.param_exprs[name]: expr for name, expr in with_defaults.items()}
+        )
+        return substituted_expr.col_type
+
     def _docstring(self) -> Optional[str]:
         if isinstance(self.templates[0].expr, exprs.FunctionCall):
             return self.templates[0].expr.fn._docstring()
@@ -97,6 +113,10 @@ class ExprTemplateFunction(Function):
 
     @property
     def display_name(self) -> str:
+        if not self.self_name and isinstance(self.templates[0].expr, exprs.FunctionCall):
+            # In the common case where the templated expression is itself a FunctionCall,
+            # fall back on the display name of the underlying FunctionCall
+            return self.templates[0].expr.fn.display_name
         return self.self_name
 
     @property

pixeltable/func/function.py

@@ -10,7 +10,6 @@ from typing import TYPE_CHECKING, Any, Callable, Optional, Sequence, cast
 import sqlalchemy as sql
 from typing_extensions import Self
 
-import pixeltable as pxt
 import pixeltable.exceptions as excs
 import pixeltable.type_system as ts
 
@@ -155,7 +154,7 @@ class Function(ABC):
         """
         raise NotImplementedError()
 
-    def __call__(self, *args: Any, **kwargs: Any) -> 'pxt.exprs.FunctionCall':
+    def __call__(self, *args: Any, **kwargs: Any) -> 'exprs.FunctionCall':
         from pixeltable import exprs
 
         args = [exprs.Expr.from_object(arg) for arg in args]
@@ -246,7 +245,7 @@ class Function(ABC):
         # `None` when any of its non-nullable inputs are `None`.
         for arg_name, arg in bound_args.items():
             param = self.signature.parameters[arg_name]
-            if param.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}:
+            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
                 continue
             if arg.col_type.nullable and not param.col_type.nullable:
                 return_type = return_type.copy(nullable=True)
@@ -385,10 +384,10 @@ class Function(ABC):
             else:
                 var = exprs.Variable(name, param.col_type)
                 bindings[name] = var
-                if args_ok and param.kind in {
+                if args_ok and param.kind in (
                     inspect.Parameter.POSITIONAL_ONLY,
                     inspect.Parameter.POSITIONAL_OR_KEYWORD,
-                }:
+                ):
                     template_args.append(var)
                 else:
                     template_kwargs[name] = var

pixeltable/func/function_registry.py

@@ -31,7 +31,7 @@ class FunctionRegistry:
             cls._instance = FunctionRegistry()
         return cls._instance
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.stored_fns_by_id: dict[UUID, Function] = {}
         self.module_fns: dict[str, Function] = {}  # fqn -> Function
         self.type_methods: dict[ts.ColumnType.Type, dict[str, Function]] = {}

pixeltable/func/signature.py

@@ -253,7 +253,7 @@ class Signature:
                 continue  # skip 'self' or 'cls' parameter
             if param.name in cls.SPECIAL_PARAM_NAMES:
                 raise excs.Error(f'{param.name!r} is a reserved parameter name')
-            if param.kind in {inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD}:
+            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
                 parameters.append(Parameter(param.name, col_type=None, kind=param.kind))
                 continue
 

pixeltable/func/tools.py

@@ -51,10 +51,10 @@ class Tool(pydantic.BaseModel):
     # The output of `tool_calls` must be a dict in standardized tool invocation format:
     # {tool_name: [{'args': {name1: value1, name2: value2, ...}}, ...], ...}
     def invoke(self, tool_calls: 'exprs.Expr') -> 'exprs.Expr':
-        from pixeltable import exprs
+        import pixeltable.functions as pxtf
 
         func_name = self.name or self.fn.name
-        return exprs.JsonMapper(tool_calls[func_name]['*'], self.__invoke_kwargs(exprs.RELATIVE_PATH_ROOT.args))
+        return pxtf.map(tool_calls[func_name]['*'], lambda x: self.__invoke_kwargs(x.args))
 
     def __invoke_kwargs(self, kwargs: 'exprs.Expr') -> 'exprs.FunctionCall':
         kwargs = {param.name: self.__extract_tool_arg(param, kwargs) for param in self.parameters.values()}

pixeltable/func/udf.py

@@ -43,7 +43,7 @@ def udf(
 ) -> ExprTemplateFunction: ...
 
 
-def udf(*args, **kwargs):
+def udf(*args, **kwargs):  # type: ignore[no-untyped-def]
     """A decorator to create a Function from a function definition.
 
     Examples:
@@ -79,7 +79,7 @@ def udf(*args, **kwargs):
         if len(args) > 0:
             raise excs.Error('Unexpected @udf decorator arguments.')
 
-        def decorator(decorated_fn: Callable):
+        def decorator(decorated_fn: Callable) -> CallableFunction:
             return make_function(
                 decorated_fn,
                 batch_size=batch_size,

pixeltable/functions/__init__.py

@@ -24,10 +24,10 @@ from . import (
     vision,
     whisper,
 )
-from .globals import count, max, mean, min, sum
+from .globals import count, map, max, mean, min, sum
 
 __all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/anthropic.py

@@ -39,7 +39,7 @@ def _anthropic_client() -> 'anthropic.AsyncAnthropic':
 
 
 class AnthropicRateLimitsInfo(env.RateLimitsInfo):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(self._get_request_resources)
 
     def _get_request_resources(self, messages: dict, max_tokens: int) -> dict[str, int]:
@@ -236,5 +236,5 @@ def _opt(arg: _T) -> Union[_T, 'anthropic.NotGiven']:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/audio.py

@@ -26,5 +26,5 @@ def get_metadata(audio: pxt.Audio) -> dict:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/deepseek.py

@@ -117,5 +117,5 @@ async def chat_completions(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/fireworks.py

@@ -131,5 +131,5 @@ async def chat_completions(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/globals.py

@@ -1,15 +1,14 @@
 import builtins
 import typing
-
-from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
-from typing import Optional, Union
+from typing import Any, Callable, Optional, Union
 
 import sqlalchemy as sql
 
-import pixeltable.type_system as ts
-from pixeltable import exprs, func
+from pixeltable import exceptions as excs, exprs, func, type_system as ts
 from pixeltable.utils.code import local_public_names
 
+from typing import _GenericAlias  # type: ignore[attr-defined]  # isort: skip
+
 
 # TODO: remove and replace calls with astype()
 def cast(expr: exprs.Expr, target_type: Union[ts.ColumnType, type, _GenericAlias]) -> exprs.Expr:
@@ -24,7 +23,7 @@ T = typing.TypeVar('T')
 class sum(func.Aggregator, typing.Generic[T]):
     """Sums the selected integers or floats."""
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.sum: T = None
 
     def update(self, val: T) -> None:
@@ -68,7 +67,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     ),
 )
 class count(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.count = 0
 
     def update(self, val: T) -> None:
@@ -89,7 +88,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
 )
 class min(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.val: T = None
 
     def update(self, val: T) -> None:
@@ -119,7 +118,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     type_substitutions=tuple({T: Optional[t]} for t in (str, int, float, bool, ts.Timestamp)),  # type: ignore[misc]
 )
 class max(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.val: T = None
 
     def update(self, val: T) -> None:
@@ -144,7 +143,7 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
 
 @func.uda(type_substitutions=({T: Optional[int]}, {T: Optional[float]}))  # type: ignore[misc]
 class mean(func.Aggregator, typing.Generic[T]):
-    def __init__(self):
+    def __init__(self) -> None:
         self.sum: T = None
         self.count = 0
 
@@ -168,8 +167,20 @@ def _(val: sql.ColumnElement) -> Optional[sql.ColumnElement]:
     return sql.sql.func.avg(val)
 
 
+def map(expr: exprs.Expr, fn: Callable[[exprs.Expr], Any]) -> exprs.Expr:
+    target_expr: exprs.Expr
+    try:
+        target_expr = exprs.Expr.from_object(fn(exprs.json_path.RELATIVE_PATH_ROOT))
+    except Exception as e:
+        raise excs.Error(
+            'Failed to evaluate map function. '
+            '(The `fn` argument to `map()` must produce a valid Pixeltable expression.)'
+        ) from e
+    return exprs.JsonMapper(expr, target_expr)
+
+
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/huggingface.py

@@ -493,5 +493,5 @@ _processor_cache: dict[tuple[str, Callable], Any] = {}
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/image.py

@@ -426,5 +426,5 @@ def mode(self: PIL.Image.Image) -> str:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/json.py

@@ -37,5 +37,5 @@ class make_list(pxt.Aggregator):
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/llama_cpp.py

@@ -100,5 +100,5 @@ _IS_GPU_AVAILABLE: Optional[bool] = None
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/math.py

@@ -63,5 +63,5 @@ def _(self: sql.ColumnElement, digits: Optional[sql.ColumnElement] = None) -> sq
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/mistralai.py

@@ -193,5 +193,5 @@ def _opt(arg: Optional[_T]) -> Union[_T, 'mistralai.types.basemodel.Unset']:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/ollama.py

@@ -129,5 +129,5 @@ def embed(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/openai.py

@@ -122,7 +122,7 @@ class OpenAIRateLimitsInfo(env.RateLimitsInfo):
 _header_duration_pattern = re.compile(r'(?:(\d+)d)?(?:(\d+)h)?(?:(\d+)ms)|(?:(\d+)m)?(?:([\d.]+)s)?')
 
 
-def _parse_header_duration(duration_str):
+def _parse_header_duration(duration_str: str) -> datetime.timedelta:
     match = _header_duration_pattern.match(duration_str)
     if not match:
         raise ValueError('Invalid duration format')
@@ -837,5 +837,5 @@ def _opt(arg: _T) -> Union[_T, 'openai.NotGiven']:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/replicate.py

@@ -69,5 +69,5 @@ async def run(input: dict[str, Any], *, ref: str) -> dict[str, Any]:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/string.py

@@ -675,5 +675,5 @@ def zfill(self: str, width: int) -> str:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/timestamp.py

@@ -312,5 +312,5 @@ def posix_timestamp(self: datetime) -> float:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/together.py

@@ -298,5 +298,5 @@ async def image_generations(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/util.py

@@ -21,7 +21,7 @@ def normalize_image_mode(image: PIL.Image.Image) -> PIL.Image.Image:
     Converts grayscale images to 3-channel for compatibility with models that only work with
     multichannel input.
     """
-    if image.mode in {'1', 'L'}:
+    if image.mode in ('1', 'L'):
         return image.convert('RGB')
     if image.mode == 'LA':
         return image.convert('RGBA')

pixeltable/functions/video.py

@@ -143,7 +143,7 @@ def _get_metadata(path: str) -> dict:
 
 
 def __get_stream_metadata(stream: av.stream.Stream) -> dict:
-    if stream.type not in {'audio', 'video'}:
+    if stream.type not in ('audio', 'video'):
         return {'type': stream.type}  # Currently unsupported
 
     codec_context = stream.codec_context
@@ -190,5 +190,5 @@ def __get_stream_metadata(stream: av.stream.Stream) -> dict:
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__

pixeltable/functions/vision.py

@@ -238,7 +238,7 @@ class mean_ap(pxt.Aggregator):
     - A `dict[int, float]` mapping each label class to an average precision (AP) value for that class.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.class_tpfp: dict[int, list[dict]] = defaultdict(list)
 
     def update(self, eval_dicts: list[dict]) -> None:
@@ -394,5 +394,5 @@ def draw_bounding_boxes(
 __all__ = local_public_names(__name__)
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return __all__