pixeltable 0.1.1__py3-none-any.whl → 0.2.0__py3-none-any.whl

pixeltable/functions/huggingface.py

@@ -0,0 +1,120 @@
+from typing import Any, Callable
+
+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
+
+
+@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])]
+
+
+@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')
+    from transformers import CLIPModel, CLIPProcessor
+
+    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
+
+    inputs = processor(text=text, return_tensors='pt', padding=True, truncation=True)
+    embeddings = model.get_text_features(**inputs).detach().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')
+    from transformers import CLIPModel, CLIPProcessor
+
+    model = _lookup_model(model_id, CLIPModel.from_pretrained)
+    processor = _lookup_processor(model_id, CLIPProcessor.from_pretrained)
+
+    inputs = processor(images=image, return_tensors='pt', padding=True)
+    embeddings = model.get_image_features(**inputs).detach().numpy()
+    return [embeddings[i] for i in range(embeddings.shape[0])]
+
+
+@pxt.udf(batch_size=32)
+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')
+    from transformers import DetrImageProcessor, DetrForObjectDetection
+
+    model = _lookup_model(model_id, lambda x: DetrForObjectDetection.from_pretrained(x, revision='no_timm'))
+    processor = _lookup_processor(model_id, lambda x: DetrImageProcessor.from_pretrained(x, revision='no_timm'))
+
+    inputs = processor(images=image, return_tensors='pt')
+    outputs = model(**inputs)
+
+    results = processor.post_process_object_detection(outputs, threshold=threshold)
+    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
+    ]
+
+
+def _lookup_model(model_id: str, create: Callable) -> Any:
+    key = (model_id, create)  # For safety, include the `create` callable in the cache key
+    if key not in _model_cache:
+        _model_cache[key] = create(model_id)
+    return _model_cache[key]
+
+
+def _lookup_processor(model_id: str, create: Callable) -> Any:
+    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')

pixeltable/functions/openai.py

@@ -0,0 +1,88 @@
+import base64
+import io
+from typing import Optional
+
+import PIL.Image
+import numpy as np
+
+import pixeltable as pxt
+import pixeltable.type_system as ts
+from pixeltable import env
+from pixeltable.func import Batch
+
+
+@pxt.udf
+def chat_completions(
+        messages: list,
+        model: str,
+        frequency_penalty: Optional[float] = None,
+        logit_bias: Optional[dict] = None,
+        max_tokens: Optional[int] = None,
+        n: Optional[int] = None,
+        presence_penalty: Optional[float] = None,
+        response_format: Optional[dict] = None,
+        seed: Optional[int] = None,
+        top_p: Optional[float] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    from openai._types import NOT_GIVEN
+    result = env.Env.get().openai_client.chat.completions.create(
+        messages=messages,
+        model=model,
+        frequency_penalty=frequency_penalty if frequency_penalty is not None else NOT_GIVEN,
+        logit_bias=logit_bias if logit_bias is not None else NOT_GIVEN,
+        max_tokens=max_tokens if max_tokens is not None else NOT_GIVEN,
+        n=n if n is not None else NOT_GIVEN,
+        presence_penalty=presence_penalty if presence_penalty is not None else NOT_GIVEN,
+        response_format=response_format if response_format is not None else NOT_GIVEN,
+        seed=seed if seed is not None else NOT_GIVEN,
+        top_p=top_p if top_p is not None else NOT_GIVEN,
+        temperature=temperature if temperature is not None else NOT_GIVEN
+    )
+    return result.dict()
+
+
+@pxt.udf
+def vision(
+        prompt: str,
+        image: PIL.Image.Image,
+        model: str = 'gpt-4-vision-preview'
+) -> str:
+    bytes_arr = io.BytesIO()
+    image.save(bytes_arr, format='png')
+    b64_bytes = base64.b64encode(bytes_arr.getvalue())
+    b64_encoded_image = b64_bytes.decode('utf-8')
+    messages = [
+        {'role': 'user',
+         'content': [
+             {'type': 'text', 'text': prompt},
+             {'type': 'image_url', 'image_url': {
+                 'url': f'data:image/png;base64,{b64_encoded_image}'
+             }}
+         ]}
+    ]
+    result = env.Env.get().openai_client.chat.completions.create(
+        messages=messages,
+        model=model
+    )
+    return result.choices[0].message.content
+
+
+@pxt.udf
+def moderations(input: str, model: Optional[str] = None) -> dict:
+    result = env.Env().get().openai_client.moderations.create(input=input, model=model)
+    return result.dict()
+
+
+@pxt.udf(batch_size=32, return_type=ts.ArrayType((None,), dtype=ts.FloatType()))
+def embeddings(input: Batch[str], *, model: str) -> Batch[np.ndarray]:
+    result = env.Env().get().openai_client.embeddings.create(
+        input=input,
+        model=model,
+        encoding_format='float'
+    )
+    embeddings = [
+        np.array(data.embedding, dtype=np.float64)
+        for data in result.data
+    ]
+    return embeddings

pixeltable/functions/pil/image.py

@@ -1,9 +1,150 @@
-from pixeltable.type_system import FloatType, ImageType
-from pixeltable.function import Function
+from typing import Dict, Any, Tuple, Optional
 
+import PIL.Image
 
-alpha_composite = Function(
-    ImageType(), [ImageType(), ImageType()], module_name='PIL.Image', eval_symbol='alpha_composite')
-blend = Function(ImageType(), [ImageType(), ImageType(), FloatType()], module_name='PIL.Image', eval_symbol='blend')
-composite = Function(
-    ImageType(), [ImageType(), ImageType(), ImageType()], module_name='PIL.Image', eval_symbol='composite')
+from pixeltable.type_system import FloatType, ImageType, IntType, ArrayType, ColumnType, StringType, JsonType, BoolType
+import pixeltable.func as func
+
+
+def _caller_return_type(bound_args: Optional[Dict[str, Any]]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    return bound_args['self'].col_type
+
+@func.udf(
+    py_fn=PIL.Image.alpha_composite, return_type=ImageType(), param_types=[ImageType(), ImageType()])
+def alpha_composite(im1: PIL.Image.Image, im2: PIL.Image.Image) -> PIL.Image.Image:
+    pass
+@func.udf(
+    py_fn=PIL.Image.blend, return_type=ImageType(), param_types=[ImageType(), ImageType(), FloatType()])
+def blend(im1: PIL.Image.Image, im2: PIL.Image.Image, alpha: float) -> PIL.Image.Image:
+    pass
+@func.udf(
+    py_fn=PIL.Image.composite, return_type=ImageType(), param_types=[ImageType(), ImageType(), ImageType()])
+def composite(image1: PIL.Image.Image, image2: PIL.Image.Image, mask: PIL.Image.Image) -> PIL.Image.Image:
+    pass
+
+
+# PIL.Image.Image methods
+
+# Image.convert()
+def _convert_return_type(bound_args: Dict[str, Any]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    assert 'self' in bound_args
+    assert 'mode' in bound_args
+    img_type = bound_args['self'].col_type
+    return ImageType(size=img_type.size, mode=bound_args['mode'])
+@func.udf(return_type=_convert_return_type, param_types=[ImageType(), StringType()])
+def convert(self: PIL.Image.Image, mode: str) -> PIL.Image.Image:
+    return self.convert(mode)
+
+# Image.crop()
+def _crop_return_type(bound_args: Dict[str, Any]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    img_type = bound_args['self'].col_type
+    box = bound_args['box']
+    if isinstance(box, list) and all(isinstance(x, int) for x in box):
+        return ImageType(size=(box[2] - box[0], box[3] - box[1]), mode=img_type.mode)
+    return ImageType()  # we can't compute the size statically
+@func.udf(
+    py_fn=PIL.Image.Image.crop, return_type=_crop_return_type,
+    param_types=[ImageType(), ArrayType((4,), dtype=IntType())])
+def crop(self: PIL.Image.Image, box: Tuple[int, int, int, int]) -> PIL.Image.Image:
+    pass
+
+# Image.getchannel()
+def _getchannel_return_type(bound_args: Dict[str, Any]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    img_type = bound_args['self'].col_type
+    return ImageType(size=img_type.size, mode='L')
+@func.udf(
+    py_fn=PIL.Image.Image.getchannel, return_type=_getchannel_return_type, param_types=[ImageType(), IntType()])
+def getchannel(self: PIL.Image.Image, channel: int) -> PIL.Image.Image:
+    pass
+
+# Image.resize()
+def resize_return_type(bound_args: Dict[str, Any]) -> ColumnType:
+    if bound_args is None:
+        return ImageType()
+    assert 'size' in bound_args
+    return ImageType(size=bound_args['size'])
+@func.udf(return_type=resize_return_type, param_types=[ImageType(), ArrayType((2, ), dtype=IntType())])
+def resize(self: PIL.Image.Image, size: Tuple[int, int]) -> PIL.Image.Image:
+    return self.resize(size)
+
+# Image.rotate()
+@func.udf(return_type=ImageType(), param_types=[ImageType(), IntType()])
+def rotate(self: PIL.Image.Image, angle: int) -> PIL.Image.Image:
+    return self.rotate(angle)
+
+# Image.transform()
+@func.udf(return_type= _caller_return_type, param_types=[ImageType(), ArrayType((2,), dtype=IntType()), IntType()])
+def transform(self: PIL.Image.Image, size: Tuple[int, int], method: int) -> PIL.Image.Image:
+    return self.transform(size, method)
+
+@func.udf(
+    py_fn=PIL.Image.Image.effect_spread, return_type=_caller_return_type, param_types=[ImageType(), FloatType()])
+def effect_spread(self: PIL.Image.Image, distance: float) -> PIL.Image.Image:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.entropy, return_type=FloatType(), param_types=[ImageType(), ImageType(), JsonType()])
+def entropy(self: PIL.Image.Image, mask: PIL.Image.Image, histogram: Dict) -> float:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getbands, return_type=JsonType(), param_types=[ImageType()])
+def getbands(self: PIL.Image.Image) -> Tuple[str]:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getbbox, return_type=JsonType(), param_types=[ImageType()])
+def getbbox(self: PIL.Image.Image) -> Tuple[int, int, int, int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getcolors, return_type=JsonType(), param_types=[ImageType(), IntType()])
+def getcolors(self: PIL.Image.Image, maxcolors: int) -> Tuple[Tuple[int, int, int], int]:
+    pass
+
+@func.udf(py_fn=PIL.Image.Image.getextrema, return_type=JsonType(), param_types=[ImageType()])
+def getextrema(self: PIL.Image.Image) -> Tuple[int, int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getpalette, return_type=JsonType(), param_types=[ImageType(), StringType()])
+def getpalette(self: PIL.Image.Image, mode: str) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getpixel, return_type=JsonType(), param_types=[ImageType(), ArrayType((2,), dtype=IntType())])
+def getpixel(self: PIL.Image.Image, xy: Tuple[int, int]) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.getprojection, return_type=JsonType(), param_types=[ImageType()])
+def getprojection(self: PIL.Image.Image) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.histogram, return_type=JsonType(), param_types=[ImageType(), ImageType(), JsonType()])
+def histogram(self: PIL.Image.Image, mask: PIL.Image.Image, histogram: Dict) -> Tuple[int]:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.quantize, return_type=ImageType(),
+    param_types=[ImageType(), IntType(), IntType(nullable=True), IntType(), IntType(nullable=True), IntType()])
+def quantize(
+        self: PIL.Image.Image, colors: int, method: int, kmeans: int, palette: int, dither: int) -> PIL.Image.Image:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.reduce, return_type=ImageType(), param_types=[ImageType(), IntType(), JsonType()])
+def reduce(self: PIL.Image.Image, factor: int, filter: Tuple[int]) -> PIL.Image.Image:
+    pass
+
+@func.udf(
+    py_fn=PIL.Image.Image.transpose, return_type=_caller_return_type, param_types=[ImageType(), IntType()])
+def transpose(self: PIL.Image.Image, method: int) -> PIL.Image.Image:
+    pass

pixeltable/functions/string.py

@@ -0,0 +1,13 @@
+from typing import Any
+
+from pixeltable.type_system import StringType
+import pixeltable.func as func
+
+
+@func.udf(return_type=StringType(), param_types=[StringType()])
+def str_format(format_str: str, *args: Any, **kwargs: Any) -> str:
+    """ Return a formatted version of format_str, using substitutions from args and kwargs:
+    - {<int>} will be replaced by the corresponding element in args
+    - {<key>} will be replaced by the corresponding value in kwargs
+    """
+    return format_str.format(*args, **kwargs)

pixeltable/functions/together.py

@@ -0,0 +1,27 @@
+from typing import Optional
+
+import pixeltable as pxt
+
+
+@pxt.udf
+def completions(
+        prompt: str,
+        model: str,
+        max_tokens: Optional[int] = None,
+        repetition_penalty: Optional[float] = None,
+        stop: Optional[list] = None,
+        top_k: Optional[int] = None,
+        top_p: Optional[float] = None,
+        temperature: Optional[float] = None
+) -> dict:
+    import together
+    return together.Complete.create(
+        prompt,
+        model,
+        max_tokens=max_tokens,
+        repetition_penalty=repetition_penalty,
+        stop=stop,
+        top_k=top_k,
+        top_p=top_p,
+        temperature=temperature
+    )

pixeltable/functions/util.py

@@ -0,0 +1,41 @@
+from typing import Tuple, List, Optional
+import types
+import sys
+
+import pixeltable.func as func
+import pixeltable.type_system as ts
+import pixeltable.env as env
+
+
+def create_nos_modules() -> List[types.ModuleType]:
+    """Create module pixeltable.functions.nos with one submodule per task and return the submodules"""
+    models = env.Env.get().nos_client.ListModels()
+    model_info = [env.Env.get().nos_client.GetModelInfo(model) for model in models]
+    model_info.sort(key=lambda info: info.task.value)
+
+    module_name = 'pixeltable.functions.nos'
+    nos_module = types.ModuleType(module_name)
+    nos_module.__package__ = 'pixeltable.functions'
+    sys.modules[module_name] = nos_module
+
+    prev_task = ''
+    new_modules: List[types.ModuleType] = []
+    sub_module: Optional[types.ModuleType] = None
+    for info in model_info:
+        if info.task.value != prev_task:
+            # we construct one submodule per task
+            namespace = info.task.name.lower()
+            submodule_name = f'{module_name}.{namespace}'
+            sub_module = types.ModuleType(submodule_name)
+            sub_module.__package__ = module_name
+            setattr(nos_module, namespace, sub_module)
+            new_modules.append(sub_module)
+            sys.modules[submodule_name] = sub_module
+            prev_task = info.task.value
+
+        # add a Function for this model to the module
+        model_id = info.name.replace("/", "_").replace("-", "_")
+        pt_func = func.NOSFunction(info, f'{submodule_name}.{model_id}')
+        setattr(sub_module, model_id, pt_func)
+
+    return new_modules

pixeltable/functions/video.py

@@ -0,0 +1,62 @@
+from typing import Optional
+import uuid
+import av
+import sys
+
+import pixeltable.env as env
+import pixeltable.func as func
+import pixeltable.type_system as ts
+
+
+_format_defaults = { # format -> (codec, ext)
+    'wav': ('pcm_s16le', 'wav'),
+    'mp3': ('libmp3lame', 'mp3'),
+    'flac': ('flac', 'flac'),
+    #'mp4': ('aac', 'm4a'),
+}
+
+# for mp4:
+# - extract_audio() fails with "Application provided invalid, non monotonically increasing dts to muxer in stream 0: 1146 >= 290"
+# - chatgpt suggests this can be fixed in the following manner
+#     for packet in container.demux(audio_stream):
+#         packet.pts = None  # Reset the PTS and DTS to allow FFmpeg to set them automatically
+#         packet.dts = None
+#         for frame in packet.decode():
+#             frame.pts = None
+#             for packet in output_stream.encode(frame):
+#                 output_container.mux(packet)
+#
+#     # Flush remaining packets
+#     for packet in output_stream.encode():
+#         output_container.mux(packet)
+
+
+_extract_audio_param_types = [
+    ts.VideoType(nullable=False),
+    ts.IntType(nullable=False),
+    ts.StringType(nullable=False),
+    ts.StringType(nullable=False)
+]
+@func.udf(return_type=ts.AudioType(nullable=True), param_types=_extract_audio_param_types)
+def extract_audio(
+        video_path: str, stream_idx: int = 0, format: str = 'wav', codec: Optional[str] = None
+) -> Optional[str]:
+    """Extract an audio stream from a video file, save it as a media file and return its path"""
+    if format not in _format_defaults:
+        raise ValueError(f'extract_audio(): unsupported audio format: {format}')
+    default_codec, ext = _format_defaults[format]
+
+    with av.open(video_path) as container:
+        if len(container.streams.audio) <= stream_idx:
+            return None
+        audio_stream = container.streams.audio[stream_idx]
+        # create this in our tmp directory, so it'll get cleaned up if it's being generated as part of a query
+        output_filename = str(env.Env.get().tmp_dir / f"{uuid.uuid4()}.{ext}")
+
+        with av.open(output_filename, "w", format=format) as output_container:
+            output_stream = output_container.add_stream(codec or default_codec)
+            for packet in container.demux(audio_stream):
+                for frame in packet.decode():
+                    output_container.mux(output_stream.encode(frame))
+
+        return output_filename

pixeltable/iterators/__init__.py

@@ -0,0 +1,3 @@
+from .base import ComponentIterator
+from .video import FrameIterator
+

pixeltable/iterators/base.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+from typing import Dict, Any, Tuple, List
+from abc import abstractmethod, ABC
+
+from pixeltable.type_system import ColumnType
+
+
+class ComponentIterator(ABC):
+    """Base class for iterators."""
+
+    @classmethod
+    @abstractmethod
+    def input_schema(cls) -> Dict[str, ColumnType]:
+        """Provide the Pixeltable types of the init() parameters
+
+        The keys need to match the names of the init() parameters. This is equivalent to the parameters_types
+        parameter of the @function decorator.
+        """
+        raise NotImplementedError
+
+    @classmethod
+    @abstractmethod
+    def output_schema(cls, *args: Any, **kwargs: Any) -> Tuple[Dict[str, ColumnType], List[str]]:
+        """Specify the dictionary returned by next() and a list of unstored column names
+
+        Returns:
+            a dictionary which is turned into a list of columns in the output table
+            a list of unstored column names
+        """
+        raise NotImplementedError
+
+    def __iter__(self) -> ComponentIterator:
+        return self
+
+    @abstractmethod
+    def __next__(self) -> Dict[str, Any]:
+        """Return the next element of the iterator as a dictionary or raise StopIteration"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def close(self) -> None:
+        """Close the iterator and release all resources"""
+        raise NotImplementedError
+
+    @abstractmethod
+    def set_pos(self, pos: int) -> None:
+        """Set the iterator position to pos"""
+        raise NotImplementedError