pixeltable 0.0.0__py3-none-any.whl

pixeltable/func/udf.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import inspect
+from typing import List, Callable, Optional, overload, Any
+
+import pixeltable as pxt
+import pixeltable.exceptions as excs
+import pixeltable.type_system as ts
+from .callable_function import CallableFunction
+from .expr_template_function import ExprTemplateFunction
+from .function import Function
+from .function_registry import FunctionRegistry
+from .globals import validate_symbol_path
+from .signature import Signature
+
+
+# Decorator invoked without parentheses: @pxt.udf
+@overload
+def udf(decorated_fn: Callable) -> Function: ...
+
+
+# Decorator schema invoked with parentheses: @pxt.udf(**kwargs)
+@overload
+def udf(
+        *,
+        return_type: Optional[ts.ColumnType] = None,
+        param_types: Optional[List[ts.ColumnType]] = None,
+        batch_size: Optional[int] = None,
+        substitute_fn: Optional[Callable] = None,
+        _force_stored: bool = False
+) -> Callable: ...
+
+
+def udf(*args, **kwargs):
+    """A decorator to create a Function from a function definition.
+
+    Examples:
+        >>> @pxt.udf
+        ... def my_function(x: int) -> int:
+        ...    return x + 1
+
+        >>> @pxt.udf(param_types=[pxt.IntType()], return_type=pxt.IntType())
+        ... def my_function(x):
+        ...    return x + 1
+    """
+    if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
+
+        # Decorator invoked without parentheses: @pxt.udf
+        # Simply call make_function with defaults.
+        return make_function(decorated_fn=args[0])
+
+    else:
+
+        # Decorator schema invoked with parentheses: @pxt.udf(**kwargs)
+        # Create a decorator for the specified schema.
+        return_type = kwargs.pop('return_type', None)
+        param_types = kwargs.pop('param_types', None)
+        batch_size = kwargs.pop('batch_size', None)
+        substitute_fn = kwargs.pop('py_fn', None)
+        force_stored = kwargs.pop('_force_stored', False)
+
+        def decorator(decorated_fn: Callable):
+            return make_function(
+                decorated_fn, return_type, param_types, batch_size,
+                substitute_fn=substitute_fn, force_stored=force_stored)
+
+        return decorator
+
+
+def make_function(
+    decorated_fn: Callable,
+    return_type: Optional[ts.ColumnType] = None,
+    param_types: Optional[List[ts.ColumnType]] = None,
+    batch_size: Optional[int] = None,
+    substitute_fn: Optional[Callable] = None,
+    function_name: Optional[str] = None,
+    force_stored: bool = False
+) -> Function:
+    """
+    Constructs a `CallableFunction` from the specified parameters.
+    If `substitute_fn` is specified, then `decorated_fn`
+    will be used only for its signature, with execution delegated to
+    `substitute_fn`.
+    """
+    # Obtain function_path from decorated_fn when appropriate
+    if force_stored:
+        # force storing the function in the db
+        function_path = None
+    elif decorated_fn.__module__ != '__main__' and decorated_fn.__name__.isidentifier():
+        function_path = f'{decorated_fn.__module__}.{decorated_fn.__qualname__}'
+    else:
+        function_path = None
+
+    # Derive function_name, if not specified explicitly
+    if function_name is None:
+        function_name = decorated_fn.__name__
+
+    # Display name to use for error messages
+    errmsg_name = function_name if function_path is None else function_path
+
+    sig = Signature.create(decorated_fn, param_types, return_type)
+
+    # batched functions must have a batched return type
+    # TODO: remove 'Python' from the error messages when we have full inference with Annotated types
+    if batch_size is not None and not sig.is_batched:
+        raise excs.Error(f'{errmsg_name}(): batch_size is specified; Python return type must be a `Batch`')
+    if batch_size is not None and len(sig.batched_parameters) == 0:
+        raise excs.Error(f'{errmsg_name}(): batch_size is specified; at least one Python parameter must be `Batch`')
+    if batch_size is None and len(sig.batched_parameters) > 0:
+        raise excs.Error(f'{errmsg_name}(): batched parameters in udf, but no `batch_size` given')
+
+    if substitute_fn is None:
+        py_fn = decorated_fn
+    else:
+        if function_path is None:
+            raise excs.Error(f'{errmsg_name}(): @udf decorator with a `substitute_fn` can only be used in a module')
+        py_fn = substitute_fn
+
+    result = CallableFunction(
+        signature=sig, py_fn=py_fn, self_path=function_path, self_name=function_name, batch_size=batch_size)
+
+    # If this function is part of a module, register it
+    if function_path is not None:
+        # do the validation at the very end, so it's easier to write tests for other failure scenarios
+        validate_symbol_path(function_path)
+        FunctionRegistry.get().register_function(function_path, result)
+
+    return result
+
+@overload
+def expr_udf(py_fn: Callable) -> ExprTemplateFunction: ...
+
+@overload
+def expr_udf(*, param_types: Optional[List[ts.ColumnType]] = None) -> Callable: ...
+
+def expr_udf(*args: Any, **kwargs: Any) -> Any:
+    def decorator(py_fn: Callable, param_types: Optional[List[ts.ColumnType]]) -> ExprTemplateFunction:
+        if py_fn.__module__ != '__main__' and py_fn.__name__.isidentifier():
+            # this is a named function in a module
+            function_path = f'{py_fn.__module__}.{py_fn.__qualname__}'
+        else:
+            function_path = None
+
+        # TODO: verify that the inferred return type matches that of the template
+        # TODO: verify that the signature doesn't contain batched parameters
+
+        # construct Parameters from the function signature
+        params = Signature.create_parameters(py_fn, param_types=param_types)
+        import pixeltable.exprs as exprs
+        var_exprs = [exprs.Variable(param.name, param.col_type) for param in params]
+        # call the function with the parameter expressions to construct an Expr with parameters
+        template = py_fn(*var_exprs)
+        assert isinstance(template, exprs.Expr)
+        py_sig = inspect.signature(py_fn)
+        if function_path is not None:
+            validate_symbol_path(function_path)
+        return ExprTemplateFunction(template, py_signature=py_sig, self_path=function_path, name=py_fn.__name__)
+
+    if len(args) == 1:
+        assert len(kwargs) == 0 and callable(args[0])
+        return decorator(args[0], None)
+    else:
+        assert len(args) == 0 and len(kwargs) == 1 and 'param_types' in kwargs
+        return lambda py_fn: decorator(py_fn, kwargs['param_types'])

pixeltable/functions/__init__.py

@@ -0,0 +1,95 @@
+import tempfile
+from pathlib import Path
+from typing import Optional, Union
+
+import PIL.Image
+import av
+import av.container
+import av.stream
+import numpy as np
+
+import pixeltable.env as env
+import pixeltable.func as func
+# import all standard function modules here so they get registered with the FunctionRegistry
+import pixeltable.functions.pil.image
+from pixeltable import exprs
+from pixeltable.type_system import IntType, ColumnType, FloatType, ImageType, VideoType
+# automatically import all submodules so that the udfs get registered
+from . import image, string, video, huggingface
+
+# TODO: remove and replace calls with astype()
+def cast(expr: exprs.Expr, target_type: ColumnType) -> exprs.Expr:
+    expr.col_type = target_type
+    return expr
+
+@func.uda(
+    update_types=[IntType()], value_type=IntType(), allows_window=True, requires_order_by=False)
+class sum(func.Aggregator):
+    def __init__(self):
+        self.sum: Union[int, float] = 0
+    def update(self, val: Union[int, float]) -> None:
+        if val is not None:
+            self.sum += val
+    def value(self) -> Union[int, float]:
+        return self.sum
+
+
+@func.uda(
+    update_types=[IntType()], value_type=IntType(), allows_window = True, requires_order_by = False)
+class count(func.Aggregator):
+    def __init__(self):
+        self.count = 0
+    def update(self, val: int) -> None:
+        if val is not None:
+            self.count += 1
+    def value(self) -> int:
+        return self.count
+
+
+@func.uda(
+    update_types=[IntType()], value_type=FloatType(), allows_window=False, requires_order_by=False)
+class mean(func.Aggregator):
+    def __init__(self):
+        self.sum = 0
+        self.count = 0
+    def update(self, val: int) -> None:
+        if val is not None:
+            self.sum += val
+            self.count += 1
+    def value(self) -> float:
+        if self.count == 0:
+            return None
+        return self.sum / self.count
+
+
+@func.uda(
+    init_types=[IntType()], update_types=[ImageType()], value_type=VideoType(),
+    requires_order_by=True, allows_window=False)
+class make_video(func.Aggregator):
+    def __init__(self, fps: int = 25):
+        """follows https://pyav.org/docs/develop/cookbook/numpy.html#generating-video"""
+        self.container: Optional[av.container.OutputContainer] = None
+        self.stream: Optional[av.stream.Stream] = None
+        self.fps = fps
+
+    def update(self, frame: PIL.Image.Image) -> None:
+        if frame is None:
+            return
+        if self.container is None:
+            (_, output_filename) = tempfile.mkstemp(suffix='.mp4', dir=str(env.Env.get().tmp_dir))
+            self.out_file = Path(output_filename)
+            self.container = av.open(str(self.out_file), mode='w')
+            self.stream = self.container.add_stream('h264', rate=self.fps)
+            self.stream.pix_fmt = 'yuv420p'
+            self.stream.width = frame.width
+            self.stream.height = frame.height
+
+        av_frame = av.VideoFrame.from_ndarray(np.array(frame.convert('RGB')), format='rgb24')
+        for packet in self.stream.encode(av_frame):
+            self.container.mux(packet)
+
+    def value(self) -> str:
+        for packet in self.stream.encode():
+            self.container.mux(packet)
+        self.container.close()
+        return str(self.out_file)

pixeltable/functions/eval.py

@@ -0,0 +1,215 @@
+from typing import List, Tuple, Dict
+from collections import defaultdict
+import sys
+
+import numpy as np
+
+import pixeltable.type_system as ts
+import pixeltable.func as func
+
+
+# TODO: figure out a better submodule structure
+
+# the following function has been adapted from MMEval
+# (sources at https://github.com/open-mmlab/mmeval)
+# Copyright (c) OpenMMLab. All rights reserved.
+def calculate_bboxes_area(bboxes: np.ndarray) -> np.ndarray:
+    """Calculate area of bounding boxes.
+
+    Args:
+        bboxes (numpy.ndarray): The bboxes with shape (n, 4) or (4, ) in 'xyxy' format.
+     Returns:
+        numpy.ndarray: The area of bboxes.
+    """
+    bboxes_w = (bboxes[..., 2] - bboxes[..., 0])
+    bboxes_h = (bboxes[..., 3] - bboxes[..., 1])
+    areas = bboxes_w * bboxes_h
+    return areas
+
+# the following function has been adapted from MMEval
+# (sources at https://github.com/open-mmlab/mmeval)
+# Copyright (c) OpenMMLab. All rights reserved.
+def calculate_overlaps(bboxes1: np.ndarray, bboxes2: np.ndarray) -> np.ndarray:
+    """Calculate the overlap between each bbox of bboxes1 and bboxes2.
+
+    Args:
+        bboxes1 (numpy.ndarray): The bboxes with shape (n, 4) in 'xyxy' format.
+        bboxes2 (numpy.ndarray): The bboxes with shape (k, 4) in 'xyxy' format.
+    Returns:
+        numpy.ndarray: IoUs or IoFs with shape (n, k).
+    """
+    bboxes1 = bboxes1.astype(np.float32)
+    bboxes2 = bboxes2.astype(np.float32)
+    rows = bboxes1.shape[0]
+    cols = bboxes2.shape[0]
+    overlaps = np.zeros((rows, cols), dtype=np.float32)
+
+    if rows * cols == 0:
+        return overlaps
+
+    if bboxes1.shape[0] > bboxes2.shape[0]:
+        # Swap bboxes for faster calculation.
+        bboxes1, bboxes2 = bboxes2, bboxes1
+        overlaps = np.zeros((cols, rows), dtype=np.float32)
+        exchange = True
+    else:
+        exchange = False
+
+    # Calculate the bboxes area.
+    area1 = calculate_bboxes_area(bboxes1)
+    area2 = calculate_bboxes_area(bboxes2)
+    eps = np.finfo(np.float32).eps
+
+    for i in range(bboxes1.shape[0]):
+        x_start = np.maximum(bboxes1[i, 0], bboxes2[:, 0])
+        y_start = np.maximum(bboxes1[i, 1], bboxes2[:, 1])
+        x_end = np.minimum(bboxes1[i, 2], bboxes2[:, 2])
+        y_end = np.minimum(bboxes1[i, 3], bboxes2[:, 3])
+        overlap_w = np.maximum(x_end - x_start, 0)
+        overlap_h = np.maximum(y_end - y_start, 0)
+        overlap = overlap_w * overlap_h
+
+        union = area1[i] + area2 - overlap
+        union = np.maximum(union, eps)
+        overlaps[i, :] = overlap / union
+    return overlaps if not exchange else overlaps.T
+
+
+# the following function has been adapted from MMEval
+# (sources at https://github.com/open-mmlab/mmeval)
+# Copyright (c) OpenMMLab. All rights reserved.
+def calculate_image_tpfp(
+    pred_bboxes: np.ndarray, pred_scores: np.ndarray, gt_bboxes: np.ndarray, min_iou: float
+) -> Tuple[np.ndarray, np.ndarray]:
+    """Calculate the true positive and false positive on an image.
+
+    Args:
+        pred_bboxes (numpy.ndarray): Predicted bboxes of this image, with
+            shape (N, 5). The scores The predicted score of the bbox is
+            concatenated behind the predicted bbox.
+        gt_bboxes (numpy.ndarray): Ground truth bboxes of this image, with
+            shape (M, 4).
+        min_iou (float): The IoU threshold.
+
+    Returns:
+        tuple (tp, fp):
+
+        - tp (numpy.ndarray): Shape (N,),
+          the true positive flag of each predicted bbox on this image.
+        - fp (numpy.ndarray): Shape (N,),
+          the false positive flag of each predicted bbox on this image.
+    """
+    # Step 1. Concatenate `gt_bboxes` and `ignore_gt_bboxes`, then set
+    # the `ignore_gt_flags`.
+    # all_gt_bboxes = np.concatenate((gt_bboxes, ignore_gt_bboxes))
+    # ignore_gt_flags = np.concatenate((np.zeros(
+    #     (gt_bboxes.shape[0], 1),
+    #     dtype=bool), np.ones((ignore_gt_bboxes.shape[0], 1), dtype=bool)))
+
+    # Step 2. Initialize the `tp` and `fp` arrays.
+    num_preds = pred_bboxes.shape[0]
+    tp = np.zeros(num_preds, dtype=np.int8)
+    fp = np.zeros(num_preds, dtype=np.int8)
+
+    # Step 3. If there are no gt bboxes in this image, then all pred bboxes
+    # within area range are false positives.
+    if gt_bboxes.shape[0] == 0:
+        fp[...] = 1
+        return tp, fp
+
+    # Step 4. Calculate the IoUs between the predicted bboxes and the
+    # ground truth bboxes.
+    ious = calculate_overlaps(pred_bboxes, gt_bboxes)
+    # For each pred bbox, the max iou with all gts.
+    ious_max = ious.max(axis=1)
+    # For each pred bbox, which gt overlaps most with it.
+    ious_argmax = ious.argmax(axis=1)
+    # Sort all pred bbox in descending order by scores.
+    sorted_indices = np.argsort(-pred_scores)
+
+    # Step 5. Count the `tp` and `fp` of each iou threshold and area range.
+    # The flags that gt bboxes have been matched.
+    gt_covered_flags = np.zeros(gt_bboxes.shape[0], dtype=bool)
+
+    # Count the prediction bboxes in order of decreasing score.
+    for pred_bbox_idx in sorted_indices:
+        if ious_max[pred_bbox_idx] >= min_iou:
+            matched_gt_idx = ious_argmax[pred_bbox_idx]
+            if not gt_covered_flags[matched_gt_idx]:
+                tp[pred_bbox_idx] = 1
+                gt_covered_flags[matched_gt_idx] = True
+            else:
+                # This gt bbox has been matched and counted as fp.
+                fp[pred_bbox_idx] = 1
+        else:
+            fp[pred_bbox_idx] = 1
+
+    return tp, fp
+
+@func.udf(
+    return_type=ts.JsonType(nullable=False),
+    param_types=[
+        ts.JsonType(nullable=False),
+        ts.JsonType(nullable=False),
+        ts.JsonType(nullable=False),
+        ts.JsonType(nullable=False),
+        ts.JsonType(nullable=False)
+    ])
+def eval_detections(
+        pred_bboxes: List[List[int]], pred_labels: List[int], pred_scores: List[float],
+        gt_bboxes: List[List[int]], gt_labels: List[int]
+) -> Dict:
+    class_idxs = list(set(pred_labels + gt_labels))
+    result: List[Dict] = []
+    pred_bboxes_arr = np.asarray(pred_bboxes)
+    pred_classes_arr = np.asarray(pred_labels)
+    pred_scores_arr = np.asarray(pred_scores)
+    gt_bboxes_arr = np.asarray(gt_bboxes)
+    gt_classes_arr = np.asarray(gt_labels)
+    for class_idx in class_idxs:
+        pred_filter = pred_classes_arr == class_idx
+        gt_filter = gt_classes_arr == class_idx
+        class_pred_scores = pred_scores_arr[pred_filter]
+        tp, fp = calculate_image_tpfp(
+            pred_bboxes_arr[pred_filter], class_pred_scores, gt_bboxes_arr[gt_filter], [0.5])
+        ordered_class_pred_scores = -np.sort(-class_pred_scores)
+        result.append({
+            'min_iou': 0.5, 'class': class_idx, 'tp': tp.tolist(), 'fp': fp.tolist(),
+            'scores': ordered_class_pred_scores.tolist(), 'num_gts': gt_filter.sum().item(),
+        })
+    return result
+
+@func.uda(
+    update_types=[ts.JsonType()], value_type=ts.JsonType(), allows_std_agg=True, allows_window=False)
+class mean_ap(func.Aggregator):
+    def __init__(self):
+        self.class_tpfp: Dict[int, List[Dict]] = defaultdict(list)
+
+    def update(self, eval_dicts: List[Dict]) -> None:
+        for eval_dict in eval_dicts:
+            class_idx = eval_dict['class']
+            self.class_tpfp[class_idx].append(eval_dict)
+
+    def value(self) -> Dict:
+        eps = np.finfo(np.float32).eps
+        result: Dict[int, float] = {}
+        for class_idx, tpfp in self.class_tpfp.items():
+            a1 = [x['tp'] for x in tpfp]
+            tp = np.concatenate([x['tp'] for x in tpfp], axis=0)
+            fp = np.concatenate([x['fp'] for x in tpfp], axis=0)
+            num_gts = np.sum([x['num_gts'] for x in tpfp])
+            scores = np.concatenate([np.asarray(x['scores']) for x in tpfp])
+            sorted_idxs = np.argsort(-scores)
+            tp_cumsum = tp[sorted_idxs].cumsum()
+            fp_cumsum = fp[sorted_idxs].cumsum()
+            precision = tp_cumsum / np.maximum(tp_cumsum + fp_cumsum, eps)
+            recall = tp_cumsum / np.maximum(num_gts, eps)
+
+            mrec = np.hstack((0, recall, 1))
+            mpre = np.hstack((0, precision, 0))
+            for i in range(mpre.shape[0] - 1, 0, -1):
+                mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
+            ind = np.where(mrec[1:] != mrec[:-1])[0]
+            ap = np.sum((mrec[ind + 1] - mrec[ind]) * mpre[ind + 1])
+            result[class_idx] = ap.item()
+        return result

pixeltable/functions/fireworks.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+import fireworks.client
+
+import pixeltable as pxt
+from pixeltable import env
+
+
+def fireworks_client() -> fireworks.client.Fireworks:
+    return env.Env.get().get_client('fireworks', lambda api_key: fireworks.client.Fireworks(api_key=api_key))
+
+
+@pxt.udf
+def chat_completions(
+        messages: list[dict[str, str]],
+        *,
+        model: str,
+        max_tokens: Optional[int] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    kwargs = {
+        'max_tokens': max_tokens,
+        'top_k': top_k,
+        'top_p': top_p,
+        'temperature': temperature
+    }
+    kwargs_not_none = dict(filter(lambda x: x[1] is not None, kwargs.items()))
+    return fireworks_client().chat.completions.create(
+        model=model,
+        messages=messages,
+        **kwargs_not_none
+    ).dict()

pixeltable/functions/huggingface.py

@@ -0,0 +1,167 @@
+from typing import Callable, TypeVar, Optional
+
+import PIL.Image
+import numpy as np
+
+import pixeltable as pxt
+import pixeltable.env as env
+import pixeltable.type_system as ts
+from pixeltable.func import Batch
+from pixeltable.functions.util import resolve_torch_device
+
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
+def sentence_transformer(
+        sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False
+) -> Batch[np.ndarray]:
+    env.Env.get().require_package('sentence_transformers')
+    from sentence_transformers import SentenceTransformer
+
+    model = _lookup_model(model_id, SentenceTransformer)
+
+    array = model.encode(sentences, normalize_embeddings=normalize_embeddings)
+    return [array[i] for i in range(array.shape[0])]
+
+
+@sentence_transformer.conditional_return_type
+def _(model_id: str) -> ts.ArrayType:
+    try:
+        from sentence_transformers import SentenceTransformer
+        model = _lookup_model(model_id, SentenceTransformer)
+        return ts.ArrayType((model.get_sentence_embedding_dimension(),), dtype=ts.FloatType(), nullable=False)
+    except ImportError:
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+
+
+@pxt.udf
+def sentence_transformer_list(sentences: list, *, model_id: str, normalize_embeddings: bool = False) -> list:
+    env.Env.get().require_package('sentence_transformers')
+    from sentence_transformers import SentenceTransformer
+
+    model = _lookup_model(model_id, SentenceTransformer)
+
+    array = model.encode(sentences, normalize_embeddings=normalize_embeddings)
+    return [array[i].tolist() for i in range(array.shape[0])]
+
+
+@pxt.udf(batch_size=32)
+def cross_encoder(sentences1: Batch[str], sentences2: Batch[str], *, model_id: str) -> Batch[float]:
+    env.Env.get().require_package('sentence_transformers')
+    from sentence_transformers import CrossEncoder
+
+    model = _lookup_model(model_id, CrossEncoder)
+
+    array = model.predict([[s1, s2] for s1, s2 in zip(sentences1, sentences2)], convert_to_numpy=True)
+    return array.tolist()
+
+
+@pxt.udf
+def cross_encoder_list(sentence1: str, sentences2: list, *, model_id: str) -> list:
+    env.Env.get().require_package('sentence_transformers')
+    from sentence_transformers import CrossEncoder
+
+    model = _lookup_model(model_id, CrossEncoder)
+
+    array = model.predict([[sentence1, s2] for s2 in sentences2], convert_to_numpy=True)
+    return array.tolist()
+
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
+def clip_text(text: Batch[str], *, model_id: str) -> Batch[np.ndarray]:
+    env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
+    from transformers import CLIPModel, CLIPProcessor
+
+    model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
+    processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
+
+    with torch.no_grad():
+        inputs = processor(text=text, return_tensors='pt', padding=True, truncation=True)
+        embeddings = model.get_text_features(**inputs.to(device)).detach().to('cpu').numpy()
+
+    return [embeddings[i] for i in range(embeddings.shape[0])]
+
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False))
+def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndarray]:
+    env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
+    from transformers import CLIPModel, CLIPProcessor
+
+    model = _lookup_model(model_id, CLIPModel.from_pretrained, device=device)
+    processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
+
+    with torch.no_grad():
+        inputs = processor(images=image, return_tensors='pt', padding=True)
+        embeddings = model.get_image_features(**inputs.to(device)).detach().to('cpu').numpy()
+
+    return [embeddings[i] for i in range(embeddings.shape[0])]
+
+
+@clip_text.conditional_return_type
+@clip_image.conditional_return_type
+def _(model_id: str) -> ts.ArrayType:
+    try:
+        from transformers import CLIPModel
+        model = _lookup_model(model_id, CLIPModel.from_pretrained)
+        return ts.ArrayType((model.config.projection_dim,), dtype=ts.FloatType(), nullable=False)
+    except ImportError:
+        return ts.ArrayType((None,), dtype=ts.FloatType(), nullable=False)
+
+
+@pxt.udf(batch_size=4)
+def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5) -> Batch[dict]:
+    env.Env.get().require_package('transformers')
+    device = resolve_torch_device('auto')
+    import torch
+    from transformers import DetrImageProcessor, DetrForObjectDetection
+
+    model = _lookup_model(
+        model_id, lambda x: DetrForObjectDetection.from_pretrained(x, revision='no_timm'), device=device)
+    processor = _lookup_processor(model_id, lambda x: DetrImageProcessor.from_pretrained(x, revision='no_timm'))
+
+    with torch.no_grad():
+        inputs = processor(images=image, return_tensors='pt')
+        outputs = model(**inputs.to(device))
+        results = processor.post_process_object_detection(
+            outputs, threshold=threshold, target_sizes=[(img.height, img.width) for img in image]
+        )
+
+    return [
+        {
+            'scores': [score.item() for score in result['scores']],
+            'labels': [label.item() for label in result['labels']],
+            'label_text': [model.config.id2label[label.item()] for label in result['labels']],
+            'boxes': [box.tolist() for box in result['boxes']]
+        }
+        for result in results
+    ]
+
+
+T = TypeVar('T')
+
+
+def _lookup_model(model_id: str, create: Callable[[str], T], device: Optional[str] = None) -> T:
+    from torch import nn
+    key = (model_id, create, device)  # For safety, include the `create` callable in the cache key
+    if key not in _model_cache:
+        model = create(model_id)
+        if device is not None:
+            model.to(device)
+        if isinstance(model, nn.Module):
+            model.eval()
+        _model_cache[key] = model
+    return _model_cache[key]
+
+
+def _lookup_processor(model_id: str, create: Callable[[str], T]) -> T:
+    key = (model_id, create)  # For safety, include the `create` callable in the cache key
+    if key not in _processor_cache:
+        _processor_cache[key] = create(model_id)
+    return _processor_cache[key]
+
+
+_model_cache = {}
+_processor_cache = {}

pixeltable/functions/image.py

@@ -0,0 +1,16 @@
+import base64
+
+import PIL.Image
+
+from pixeltable.type_system import ImageType, StringType
+import pixeltable.func as func
+
+
+@func.udf
+def b64_encode(img: PIL.Image.Image, image_format: str = 'png') -> str:
+    # Encode this image as a b64-encoded png.
+    import io
+    bytes_arr = io.BytesIO()
+    img.save(bytes_arr, format=image_format)
+    b64_bytes = base64.b64encode(bytes_arr.getvalue())
+    return b64_bytes.decode('utf-8')