pixeltable 0.2.7__py3-none-any.whl → 0.2.9__py3-none-any.whl

pixeltable/functions/__init__.py

@@ -1,95 +1,10 @@
-import tempfile
-from pathlib import Path
-from typing import Optional, Union
+from . import fireworks, huggingface, image, openai, string, together, video
+from .globals import *
+from pixeltable.utils.code import local_public_names
 
-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
+__all__ = local_public_names(__name__, exclude=['globals']) + local_public_names(globals.__name__)
 
-# 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)
+def __dir__():
+    return __all__

pixeltable/functions/eval.py

@@ -10,6 +10,7 @@ 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.
@@ -21,11 +22,12 @@ def calculate_bboxes_area(bboxes: np.ndarray) -> np.ndarray:
      Returns:
         numpy.ndarray: The area of bboxes.
     """
-    bboxes_w = (bboxes[..., 2] - bboxes[..., 0])
-    bboxes_h = (bboxes[..., 3] - bboxes[..., 1])
+    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.
@@ -146,6 +148,7 @@ def calculate_image_tpfp(
 
     return tp, fp
 
+
 @func.udf(
     return_type=ts.JsonType(nullable=False),
     param_types=[
@@ -153,11 +156,15 @@ def calculate_image_tpfp(
         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]
+    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] = []
@@ -170,17 +177,22 @@ def eval_detections(
         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])
+        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(),
-        })
+        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)
+
+@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)

pixeltable/functions/fireworks.py

@@ -1,39 +1,41 @@
-from typing import Optional
-
-import fireworks.client
+from typing import Optional, TYPE_CHECKING
 
 import pixeltable as pxt
 from pixeltable import env
+from pixeltable.utils.code import local_public_names
+
+if TYPE_CHECKING:
+    import fireworks.client
 
 
 @env.register_client('fireworks')
-def _(api_key: str) -> fireworks.client.Fireworks:
+def _(api_key: str) -> 'fireworks.client.Fireworks':
+    import fireworks.client
+
     return fireworks.client.Fireworks(api_key=api_key)
 
 
-def _fireworks_client() -> fireworks.client.Fireworks:
+def _fireworks_client() -> 'fireworks.client.Fireworks':
     return env.Env.get().get_client('fireworks')
 
 
 @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
+    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 = {'max_tokens': max_tokens, 'top_k': top_k, 'top_p': top_p, 'temperature': temperature}
     kwargs_not_none = {k: v for k, v in kwargs.items() if v is not None}
-    return _fireworks_client().chat.completions.create(
-        model=model,
-        messages=messages,
-        **kwargs_not_none
-    ).dict()
+    return _fireworks_client().chat.completions.create(model=model, messages=messages, **kwargs_not_none).dict()
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/globals.py

@@ -0,0 +1,62 @@
+from typing import Union
+
+import pixeltable.func as func
+import pixeltable.type_system as ts
+from pixeltable import exprs
+from pixeltable.utils.code import local_public_names
+
+
+# TODO: remove and replace calls with astype()
+def cast(expr: exprs.Expr, target_type: ts.ColumnType) -> exprs.Expr:
+    expr.col_type = target_type
+    return expr
+
+
+@func.uda(update_types=[ts.IntType()], value_type=ts.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=[ts.IntType()], value_type=ts.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=[ts.IntType()], value_type=ts.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
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/huggingface.py

@@ -7,12 +7,13 @@ 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
+from pixeltable.functions.util import resolve_torch_device, normalize_image_mode
+from pixeltable.utils.code import local_public_names
 
 
 @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
+    sentences: Batch[str], *, model_id: str, normalize_embeddings: bool = False
 ) -> Batch[np.ndarray]:
     """Runs the specified sentence transformer model."""
     env.Env.get().require_package('sentence_transformers')
@@ -28,6 +29,7 @@ def sentence_transformer(
 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:
@@ -109,6 +111,7 @@ def clip_image(image: Batch[PIL.Image.Image], *, model_id: str) -> Batch[np.ndar
 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:
@@ -124,11 +127,13 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
     from transformers import DetrImageProcessor, DetrForObjectDetection
 
     model = _lookup_model(
-        model_id, lambda x: DetrForObjectDetection.from_pretrained(x, revision='no_timm'), device=device)
+        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'))
+    normalized_images = [normalize_image_mode(img) for img in image]
 
     with torch.no_grad():
-        inputs = processor(images=image, return_tensors='pt')
+        inputs = processor(images=normalized_images, 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]
@@ -139,7 +144,7 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
             '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']]
+            'boxes': [box.tolist() for box in result['boxes']],
         }
         for result in results
     ]
@@ -149,19 +154,10 @@ def detr_for_object_detection(image: Batch[PIL.Image.Image], *, model_id: str, t
 def detr_to_coco(image: PIL.Image.Image, detr_info: dict[str, Any]) -> dict[str, Any]:
     bboxes, labels = detr_info['boxes'], detr_info['labels']
     annotations = [
-        {
-            'bbox': [bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]],
-            'category': label
-        }
+        {'bbox': [bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]], 'category': label}
         for bbox, label in zip(bboxes, labels)
     ]
-    return {
-        'image': {
-            'width': image.width,
-            'height': image.height
-        },
-        'annotations': annotations
-    }
+    return {'image': {'width': image.width, 'height': image.height}, 'annotations': annotations}
 
 
 T = TypeVar('T')
@@ -169,6 +165,7 @@ 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)
@@ -189,3 +186,10 @@ def _lookup_processor(model_id: str, create: Callable[[str], T]) -> T:
 
 _model_cache = {}
 _processor_cache = {}
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/functions/image.py

@@ -1,16 +1,185 @@
 import base64
+from typing import Optional, Tuple
 
 import PIL.Image
 
-from pixeltable.type_system import ImageType, StringType
 import pixeltable.func as func
+import pixeltable.type_system as ts
+from pixeltable.utils.code import local_public_names
+from pixeltable.exprs import Expr
 
 
 @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')
+
+
+@func.udf(substitute_fn=PIL.Image.alpha_composite)
+def alpha_composite(im1: PIL.Image.Image, im2: PIL.Image.Image) -> PIL.Image.Image:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.blend)
+def blend(im1: PIL.Image.Image, im2: PIL.Image.Image, alpha: float) -> PIL.Image.Image:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.composite)
+def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.Image) -> PIL.Image.Image:
+    pass
+
+
+# PIL.Image.Image methods
+
+
+# Image.convert()
+@func.udf
+def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
+    return self.convert(mode)
+
+
+@convert.conditional_return_type
+def _(self: Expr, mode: str) -> ts.ColumnType:
+    input_type = self.col_type
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=input_type.size, mode=mode, nullable=input_type.nullable)
+
+
+# Image.crop()
+@func.udf(substitute_fn=PIL.Image.Image.crop, param_types=[ts.ImageType(), ts.ArrayType((4,), dtype=ts.IntType())])
+def crop(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> PIL.Image.Image:
+    pass
+
+
+@crop.conditional_return_type
+def _(self: Expr, box: Tuple[int, int, int, int]) -> ts.ColumnType:
+    input_type = self.col_type
+    assert isinstance(input_type, ts.ImageType)
+    if isinstance(box, list) and all(isinstance(x, int) for x in box):
+        return ts.ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=input_type.mode, nullable=input_type.nullable)
+    return ts.ImageType(mode=input_type.mode, nullable=input_type.nullable)  # we can't compute the size statically
+
+
+# Image.getchannel()
+@func.udf(substitute_fn=PIL.Image.Image.getchannel)
+def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
+    pass
+
+
+@getchannel.conditional_return_type
+def _(self: Expr) -> ts.ColumnType:
+    input_type = self.col_type
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=input_type.size, mode='L', nullable=input_type.nullable)
+
+
+# Image.resize()
+@func.udf(param_types=[ts.ImageType(), ts.ArrayType((2,), dtype=ts.IntType())])
+def resize(self: PIL.Image.Image, size: Tuple[int, int]) -> PIL.Image.Image:
+    return self.resize(size)
+
+
+@resize.conditional_return_type
+def _(self: Expr, size: Tuple[int, int]) -> ts.ColumnType:
+    input_type = self.col_type
+    assert isinstance(input_type, ts.ImageType)
+    return ts.ImageType(size=size, mode=input_type.mode, nullable=input_type.nullable)
+
+
+# Image.rotate()
+@func.udf
+def rotate(self: PIL.Image.Image, angle: int) -> PIL.Image.Image:
+    return self.rotate(angle)
+
+
+@func.udf(substitute_fn=PIL.Image.Image.effect_spread)
+def effect_spread(self: PIL.Image.Image, distance: int) -> PIL.Image.Image:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.transpose)
+def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
+    pass
+
+
+@rotate.conditional_return_type
+@effect_spread.conditional_return_type
+@transpose.conditional_return_type
+def _(self: Expr) -> ts.ColumnType:
+    return self.col_type
+
+
+@func.udf(substitute_fn=PIL.Image.Image.entropy)
+def entropy(self: PIL.Image.Image, mask: Optional[PIL.Image.Image] = None, extrema: Optional[list] = None) -> float:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.getbands)
+def getbands(self: PIL.Image.Image) -> Tuple[str]:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.getbbox)
+def getbbox(self: PIL.Image.Image) -> Tuple[int, int, int, int]:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.getcolors)
+def getcolors(self: PIL.Image.Image, maxcolors: int) -> Tuple[Tuple[int, int, int], int]:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.getextrema)
+def getextrema(self: PIL.Image.Image) -> Tuple[int, int]:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.getpalette)
+def getpalette(self: PIL.Image.Image, mode: Optional[str] = None) -> Tuple[int]:
+    pass
+
+
+@func.udf(param_types=[ts.ImageType(), ts.ArrayType((2,), dtype=ts.IntType())])
+def getpixel(self: PIL.Image.Image, xy: tuple[int, int]) -> Tuple[int]:
+    # `xy` will be a list; `tuple(xy)` is necessary for pillow 9 compatibility
+    return self.getpixel(tuple(xy))
+
+
+@func.udf(substitute_fn=PIL.Image.Image.getprojection)
+def getprojection(self: PIL.Image.Image) -> Tuple[int]:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.histogram)
+def histogram(self: PIL.Image.Image, mask: PIL.Image.Image, extrema: Optional[list] = None) -> Tuple[int]:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.quantize)
+def quantize(
+    self: PIL.Image.Image,
+    colors: int = 256,
+    method: Optional[int] = None,
+    kmeans: int = 0,
+    palette: Optional[int] = None,
+    dither: int = PIL.Image.Dither.FLOYDSTEINBERG,
+) -> PIL.Image.Image:
+    pass
+
+
+@func.udf(substitute_fn=PIL.Image.Image.reduce)
+def reduce(self: PIL.Image.Image, factor: int, box: Optional[Tuple[int]] = None) -> PIL.Image.Image:
+    pass
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__