pixeltable 0.3.8__py3-none-any.whl → 0.3.10__py3-none-any.whl

pixeltable/exprs/string_op.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import Any, Optional, Union
+
+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 .globals import StringOperator
+from .row_builder import RowBuilder
+from .sql_element_cache import SqlElementCache
+
+
+class StringOp(Expr):
+    """
+    Allows operations on strings
+    """
+
+    operator: StringOperator
+
+    def __init__(self, operator: StringOperator, op1: Expr, op2: Expr):
+        super().__init__(ts.StringType(nullable=op1.col_type.nullable))
+        self.operator = operator
+        self.components = [op1, op2]
+        assert op1.col_type.is_string_type()
+        if operator in {StringOperator.CONCAT, StringOperator.REPEAT}:
+            if operator == StringOperator.CONCAT and not op2.col_type.is_string_type():
+                raise excs.Error(
+                    f'{self}: {operator} on strings requires string type, but {op2} has type {op2.col_type}'
+                )
+            if operator == StringOperator.REPEAT and not op2.col_type.is_int_type():
+                raise excs.Error(f'{self}: {operator} on strings requires int type, but {op2} has type {op2.col_type}')
+        else:
+            raise excs.Error(
+                f'{self}: invalid operation {operator} on strings; '
+                f'only operators {StringOperator.CONCAT} and {StringOperator.REPEAT} are supported'
+            )
+        self.id = self._create_id()
+
+    @property
+    def _op1(self) -> Expr:
+        return self.components[0]
+
+    @property
+    def _op2(self) -> Expr:
+        return self.components[1]
+
+    def __repr__(self) -> str:
+        # add parentheses around operands that are StringOpExpr to express precedence
+        op1_str = f'({self._op1})' if isinstance(self._op1, StringOp) else str(self._op1)
+        op2_str = f'({self._op2})' if isinstance(self._op2, StringOp) else str(self._op2)
+        return f'{op1_str} {self.operator} {op2_str}'
+
+    def _equals(self, other: StringOp) -> bool:
+        return self.operator == other.operator
+
+    def _id_attrs(self) -> list[tuple[str, Any]]:
+        return [*super()._id_attrs(), ('operator', self.operator.value)]
+
+    def sql_expr(self, sql_elements: SqlElementCache) -> Optional[sql.ColumnElement]:
+        left = sql_elements.get(self._op1)
+        right = sql_elements.get(self._op2)
+        if left is None or right is None:
+            return None
+        if self.operator == StringOperator.CONCAT:
+            return left.concat(right)
+        if self.operator == StringOperator.REPEAT:
+            return sql.func.repeat(sql.cast(left, sql.String), sql.cast(right, sql.Integer))
+        return None
+
+    def eval(self, data_row: DataRow, row_builder: RowBuilder) -> None:
+        op1_val = data_row[self._op1.slot_idx]
+        op2_val = data_row[self._op2.slot_idx]
+        data_row[self.slot_idx] = self.eval_nullable(op1_val, op2_val)
+
+    def eval_nullable(self, op1_val: Union[str, None], op2_val: Union[int, str, None]) -> Union[str, None]:
+        """
+        Return the result of evaluating the expression on two nullable int/float operands,
+        None is interpreted as SQL NULL
+        """
+        if op1_val is None or op2_val is None:
+            return None
+        return self.eval_non_null(op1_val, op2_val)
+
+    def eval_non_null(self, op1_val: str, op2_val: Union[int, str]) -> str:
+        """
+        Return the result of evaluating the expression on two int/float operands
+        """
+        assert self.operator in {StringOperator.CONCAT, StringOperator.REPEAT}
+        if self.operator == StringOperator.CONCAT:
+            assert isinstance(op2_val, str)
+            return op1_val + op2_val
+        else:
+            assert isinstance(op2_val, int)
+            return op1_val * op2_val
+
+    def _as_dict(self) -> dict:
+        return {'operator': self.operator.value, **super()._as_dict()}
+
+    @classmethod
+    def _from_dict(cls, d: dict, components: list[Expr]) -> StringOp:
+        assert 'operator' in d
+        assert len(components) == 2
+        return cls(StringOperator(d['operator']), components[0], components[1])

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,47 +84,27 @@ 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__)

pixeltable/func/callable_function.py

@@ -8,6 +8,7 @@ from uuid import UUID
 import cloudpickle  # type: ignore[import-untyped]
 
 import pixeltable.exceptions as excs
+from pixeltable.utils.coroutine import run_coroutine_synchronously
 
 from .function import Function
 from .signature import Signature
@@ -93,13 +94,15 @@ class CallableFunction(Function):
             batched_kwargs = {k: [v] for k, v in kwargs.items() if k not in constant_param_names}
             result: list[Any]
             if inspect.iscoroutinefunction(self.py_fn):
-                result = asyncio.run(self.py_fn(*batched_args, **constant_kwargs, **batched_kwargs))
+                # TODO: This is temporary (see note in utils/coroutine.py)
+                result = run_coroutine_synchronously(self.py_fn(*batched_args, **constant_kwargs, **batched_kwargs))
             else:
                 result = self.py_fn(*batched_args, **constant_kwargs, **batched_kwargs)
             assert len(result) == 1
             return result[0]
         elif inspect.iscoroutinefunction(self.py_fn):
-            return asyncio.run(self.py_fn(*args, **kwargs))
+            # TODO: This is temporary (see note in utils/coroutine.py)
+            return run_coroutine_synchronously(self.py_fn(*args, **kwargs))
         else:
             return self.py_fn(*args, **kwargs)
 

pixeltable/func/query_template_function.py

@@ -17,8 +17,6 @@ class QueryTemplateFunction(Function):
 
     template_df: Optional['DataFrame']
     self_name: Optional[str]
-    # conn: Optional[sql.engine.Connection]
-    defaults: dict[str, exprs.Literal]
 
     @classmethod
     def create(
@@ -46,20 +44,6 @@ class QueryTemplateFunction(Function):
         self.self_name = name
         self.template_df = template_df
 
-        # if we're running as part of an ongoing update operation, we need to use the same connection, otherwise
-        # we end up with a deadlock
-        # TODO: figure out a more general way to make execution state available
-        # self.conn = None
-
-        # convert defaults to Literals
-        self.defaults = {}  # key: param name, value: default value converted to a Literal
-        param_types = self.template_df.parameters()
-        for param in [p for p in sig.parameters.values() if p.has_default()]:
-            assert param.name in param_types
-            param_type = param_types[param.name]
-            literal_default = exprs.Literal(param.default, col_type=param_type)
-            self.defaults[param.name] = literal_default
-
     def _update_as_overload_resolution(self, signature_idx: int) -> None:
         pass  # only one signature supported for QueryTemplateFunction
 
@@ -72,7 +56,11 @@ class QueryTemplateFunction(Function):
         bound_args = self.signature.py_signature.bind(*args, **kwargs).arguments
         # apply defaults, otherwise we might have Parameters left over
         bound_args.update(
-            {param_name: default for param_name, default in self.defaults.items() if param_name not in bound_args}
+            {
+                param.name: param.default
+                for param in self.signature.parameters.values()
+                if param.has_default() and param.name not in bound_args
+            }
         )
         bound_df = self.template_df.bind(bound_args)
         result = await bound_df._acollect()
@@ -87,7 +75,7 @@ class QueryTemplateFunction(Function):
         return self.self_name
 
     def _as_dict(self) -> dict:
-        return {'name': self.name, 'signature': self.signatures[0].as_dict(), 'df': self.template_df.as_dict()}
+        return {'name': self.name, 'signature': self.signature.as_dict(), 'df': self.template_df.as_dict()}
 
     @classmethod
     def _from_dict(cls, d: dict) -> Function:

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/functions/__init__.py

@@ -24,7 +24,7 @@ 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__)
 

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:
@@ -168,6 +167,18 @@ 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__)