pixeltable 0.2.26__py3-none-any.whl → 0.5.7__py3-none-any.whl

pixeltable/exec/expr_eval_node.py

@@ -1,232 +0,0 @@
-import logging
-import sys
-import time
-import warnings
-from dataclasses import dataclass
-from typing import Iterable, Optional
-
-from tqdm import TqdmWarning, tqdm
-
-from pixeltable import exprs
-from pixeltable.func import CallableFunction
-
-from .data_row_batch import DataRowBatch
-from .exec_node import ExecNode
-
-_logger = logging.getLogger('pixeltable')
-
-
-class ExprEvalNode(ExecNode):
-    """Materializes expressions
-    """
-    @dataclass
-    class Cohort:
-        """List of exprs that form an evaluation context and contain calls to at most one external function"""
-        exprs_: list[exprs.Expr]
-        batched_fn: Optional[CallableFunction]
-        segment_ctxs: list['exprs.RowBuilder.EvalCtx']
-        target_slot_idxs: list[int]
-        batch_size: int = 8
-
-    def __init__(
-            self, row_builder: exprs.RowBuilder, output_exprs: Iterable[exprs.Expr], input_exprs: Iterable[exprs.Expr],
-            input: ExecNode
-    ):
-        super().__init__(row_builder, output_exprs, input_exprs, input)
-        self.input_exprs = input_exprs
-        input_slot_idxs = {e.slot_idx for e in input_exprs}
-        # we're only materializing exprs that are not already in the input
-        self.target_exprs = [e for e in output_exprs if e.slot_idx not in input_slot_idxs]
-        self.pbar: Optional[tqdm] = None
-        self.cohorts: list[ExprEvalNode.Cohort] = []
-        self._create_cohorts()
-
-    def __next__(self) -> DataRowBatch:
-        input_batch = next(self.input)
-        # compute target exprs
-        for cohort in self.cohorts:
-            self._exec_cohort(cohort, input_batch)
-        _logger.debug(f'ExprEvalNode: returning {len(input_batch)} rows')
-        return input_batch
-
-    def _open(self) -> None:
-        warnings.simplefilter("ignore", category=TqdmWarning)
-        # This is a temporary hack. When B-tree indices on string columns were implemented (via computed columns
-        # that invoke the `BtreeIndex.str_filter` udf), it resulted in frivolous progress bars appearing on every
-        # insertion. This special-cases the `str_filter` call to suppress the corresponding progress bar.
-        # TODO(aaron-siegel) Remove this hack once we clean up progress bars more generally.
-        is_str_filter_node = all(
-            isinstance(expr, exprs.FunctionCall) and expr.fn.name == 'str_filter' for expr in self.output_exprs
-        )
-        if self.ctx.show_pbar and not is_str_filter_node:
-            self.pbar = tqdm(
-                total=len(self.target_exprs) * self.ctx.num_rows,
-                desc='Computing cells',
-                unit=' cells',
-                ncols=100,
-                file=sys.stdout
-            )
-
-    def _close(self) -> None:
-        if self.pbar is not None:
-            self.pbar.close()
-
-    def _get_batched_fn(self, expr: exprs.Expr) -> Optional[CallableFunction]:
-        if isinstance(expr, exprs.FunctionCall) and isinstance(expr.fn, CallableFunction) and expr.fn.is_batched:
-            return expr.fn
-        return None
-
-    def _is_batched_fn_call(self, expr: exprs.Expr) -> bool:
-        return self._get_batched_fn(expr) is not None
-
-    def _create_cohorts(self) -> None:
-        all_exprs = self.row_builder.get_dependencies(self.target_exprs)
-        # break up all_exprs into cohorts such that each cohort contains calls to at most one external function;
-        # seed the cohorts with only the ext fn calls
-        cohorts: list[list[exprs.Expr]] = []
-        current_batched_fn: Optional[CallableFunction] = None
-        for e in all_exprs:
-            if not self._is_batched_fn_call(e):
-                continue
-            assert isinstance(e, exprs.FunctionCall)
-            assert isinstance(e.fn, CallableFunction)
-            if current_batched_fn is None or current_batched_fn != e.fn:
-                # create a new cohort
-                cohorts.append([])
-                current_batched_fn = e.fn
-            cohorts[-1].append(e)
-
-        # expand the cohorts to include all exprs that are in the same evaluation context as the external calls;
-        # cohorts are evaluated in order, so we can exclude the target slots from preceding cohorts and input slots
-        exclude = set(e.slot_idx for e in self.input_exprs)
-        all_target_slot_idxs = set(e.slot_idx for e in self.target_exprs)
-        target_slot_idxs: list[list[int]] = []  # the ones materialized by each cohort
-        for i in range(len(cohorts)):
-            cohorts[i] = self.row_builder.get_dependencies(
-                cohorts[i], exclude=[self.row_builder.unique_exprs[slot_idx] for slot_idx in exclude])
-            target_slot_idxs.append(
-                [e.slot_idx for e in cohorts[i] if e.slot_idx in all_target_slot_idxs])
-            exclude.update(target_slot_idxs[-1])
-
-        all_cohort_slot_idxs = set(e.slot_idx for cohort in cohorts for e in cohort)
-        remaining_slot_idxs = set(all_target_slot_idxs) - all_cohort_slot_idxs
-        if len(remaining_slot_idxs) > 0:
-            cohorts.append(self.row_builder.get_dependencies(
-                [self.row_builder.unique_exprs[slot_idx] for slot_idx in remaining_slot_idxs],
-                exclude=[self.row_builder.unique_exprs[slot_idx] for slot_idx in exclude]))
-            target_slot_idxs.append(list(remaining_slot_idxs))
-        # we need to have captured all target slots at this point
-        assert all_target_slot_idxs == set().union(*target_slot_idxs)
-
-        for i in range(len(cohorts)):
-            cohort = cohorts[i]
-            # segment the cohort into sublists that contain either a single ext. function call or no ext. function calls
-            # (i.e., only computed cols)
-            assert len(cohort) > 0
-            # create the first segment here, so we can avoid checking for an empty list in the loop
-            segments = [[cohort[0]]]
-            is_batched_segment = self._is_batched_fn_call(cohort[0])
-            batched_fn: Optional[CallableFunction] = self._get_batched_fn(cohort[0])
-            for e in cohort[1:]:
-                if self._is_batched_fn_call(e):
-                    segments.append([e])
-                    is_batched_segment = True
-                    batched_fn = self._get_batched_fn(e)
-                else:
-                    if is_batched_segment:
-                        # start a new segment
-                        segments.append([])
-                        is_batched_segment = False
-                    segments[-1].append(e)
-
-            # we create the EvalCtxs manually because create_eval_ctx() would repeat the dependencies of each segment
-            segment_ctxs = [
-                exprs.RowBuilder.EvalCtx(
-                    slot_idxs=[e.slot_idx for e in s], exprs=s, target_slot_idxs=[], target_exprs=[])
-                for s in segments
-            ]
-            cohort_info = self.Cohort(cohort, batched_fn, segment_ctxs, target_slot_idxs[i])
-            self.cohorts.append(cohort_info)
-
-    def _exec_cohort(self, cohort: Cohort, rows: DataRowBatch) -> None:
-        """Compute the cohort for the entire input batch by dividing it up into sub-batches"""
-        batch_start_idx = 0  # start row of the current sub-batch
-        # for multi-resolution models, we re-assess the correct ext fn batch size for each input batch
-        ext_batch_size = cohort.batched_fn.get_batch_size() if cohort.batched_fn is not None else None
-        if ext_batch_size is not None:
-            cohort.batch_size = ext_batch_size
-
-        while batch_start_idx < len(rows):
-            num_batch_rows = min(cohort.batch_size, len(rows) - batch_start_idx)
-            for segment_ctx in cohort.segment_ctxs:
-                if not self._is_batched_fn_call(segment_ctx.exprs[0]):
-                    # compute batch row-wise
-                    for row_idx in range(batch_start_idx, batch_start_idx + num_batch_rows):
-                        self.row_builder.eval(
-                            rows[row_idx], segment_ctx, self.ctx.profile, ignore_errors=self.ctx.ignore_errors)
-                else:
-                    fn_call = segment_ctx.exprs[0]
-                    assert isinstance(fn_call, exprs.FunctionCall)
-                    # make a batched external function call
-                    arg_batches: list[list[exprs.Expr]] = [[] for _ in range(len(fn_call.args))]
-                    kwarg_batches: dict[str, list[exprs.Expr]] = {k: [] for k in fn_call.kwargs.keys()}
-
-                    valid_batch_idxs: list[int] = []  # rows with exceptions are not valid
-                    for row_idx in range(batch_start_idx, batch_start_idx + num_batch_rows):
-                        row = rows[row_idx]
-                        if row.has_exc(fn_call.slot_idx):
-                            # one of our inputs had an exception, skip this row
-                            continue
-                        valid_batch_idxs.append(row_idx)
-                        args, kwargs = fn_call._make_args(row)
-                        for i in range(len(args)):
-                            arg_batches[i].append(args[i])
-                        for k in kwargs.keys():
-                            kwarg_batches[k].append(kwargs[k])
-                    num_valid_batch_rows = len(valid_batch_idxs)
-
-                    if ext_batch_size is None:
-                        # we need to choose a batch size based on the args
-                        assert isinstance(fn_call.fn, CallableFunction)
-                        sample_args = [arg_batches[i][0] for i in range(len(arg_batches))]
-                        ext_batch_size = fn_call.fn.get_batch_size(*sample_args)
-
-                    num_remaining_batch_rows = num_valid_batch_rows
-                    while num_remaining_batch_rows > 0:
-                        # we make ext. fn calls in batches of ext_batch_size
-                        if ext_batch_size is None:
-                            pass
-                        num_ext_batch_rows = min(ext_batch_size, num_remaining_batch_rows)
-                        ext_batch_offset = num_valid_batch_rows - num_remaining_batch_rows  # offset into args, not rows
-                        call_args = [
-                            arg_batches[i][ext_batch_offset:ext_batch_offset + num_ext_batch_rows]
-                            for i in range(len(arg_batches))
-                        ]
-                        call_kwargs = {
-                            k: kwarg_batches[k][ext_batch_offset:ext_batch_offset + num_ext_batch_rows]
-                            for k in kwarg_batches.keys()
-                        }
-                        start_ts = time.perf_counter()
-                        assert isinstance(fn_call.fn, CallableFunction)
-                        result_batch = fn_call.fn.exec_batch(*call_args, **call_kwargs)
-                        self.ctx.profile.eval_time[fn_call.slot_idx] += time.perf_counter() - start_ts
-                        self.ctx.profile.eval_count[fn_call.slot_idx] += num_ext_batch_rows
-
-                        # move the result into the row batch
-                        for result_idx in range(len(result_batch)):
-                            row_idx = valid_batch_idxs[ext_batch_offset + result_idx]
-                            row = rows[row_idx]
-                            row[fn_call.slot_idx] = result_batch[result_idx]
-
-                        num_remaining_batch_rows -= num_ext_batch_rows
-
-                    # switch to the ext fn batch size
-                    cohort.batch_size = ext_batch_size
-
-            # make sure images for stored cols have been saved to files before moving on to the next batch
-            rows.flush_imgs(
-                slice(batch_start_idx, batch_start_idx + num_batch_rows), self.stored_img_cols, self.flushed_img_slots)
-            if self.pbar is not None:
-                self.pbar.update(num_batch_rows * len(cohort.target_slot_idxs))
-            batch_start_idx += num_batch_rows
-

pixeltable/ext/__init__.py

@@ -1,14 +0,0 @@
-"""
-Extended integrations for Pixeltable. This package contains experimental or demonstration features that
-are not intended for production use. Long-term support cannot be guaranteed, usually because the features
-have dependencies whose future support is unclear.
-"""
-
-from pixeltable.utils.code import local_public_names
-from . import functions
-
-__all__ = local_public_names(__name__)
-
-
-def __dir__():
-    return __all__

pixeltable/ext/functions/__init__.py

@@ -1,8 +0,0 @@
-from pixeltable.utils.code import local_public_names
-from . import whisperx, yolox
-
-__all__ = local_public_names(__name__)
-
-
-def __dir__():
-    return __all__

pixeltable/ext/functions/whisperx.py

@@ -1,77 +0,0 @@
-from typing import TYPE_CHECKING, Optional
-
-from pixeltable.utils.code import local_public_names
-
-if TYPE_CHECKING:
-    from whisperx.asr import FasterWhisperPipeline  # type: ignore[import-untyped]
-
-import pixeltable as pxt
-
-
-@pxt.udf
-def transcribe(
-    audio: pxt.Audio,
-    *,
-    model: str,
-    compute_type: Optional[str] = None,
-    language: Optional[str] = None,
-    chunk_size: int = 30
-) -> dict:
-    """
-    Transcribe an audio file using WhisperX.
-
-    This UDF runs a transcription model _locally_ using the WhisperX library,
-    equivalent to the WhisperX `transcribe` function, as described in the
-    [WhisperX library documentation](https://github.com/m-bain/whisperX).
-
-    WhisperX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
-
-    __Requirements:__
-
-    - `pip install whisperx`
-
-    Args:
-        audio: The audio file to transcribe.
-        model: The name of the model to use for transcription.
-
-    See the [WhisperX library documentation](https://github.com/m-bain/whisperX) for details
-    on the remaining parameters.
-
-    Returns:
-        A dictionary containing the transcription and various other metadata.
-
-    Examples:
-        Add a computed column that applies the model `tiny.en` to an existing Pixeltable column `tbl.audio`
-        of the table `tbl`:
-
-        >>> tbl['result'] = transcribe(tbl.audio, model='tiny.en')
-    """
-    import torch
-    import whisperx  # type: ignore[import-untyped]
-
-    device = 'cuda' if torch.cuda.is_available() else 'cpu'
-    compute_type = compute_type or ('float16' if device == 'cuda' else 'int8')
-    model = _lookup_model(model, device, compute_type)
-    audio_array = whisperx.load_audio(audio)
-    result = model.transcribe(audio_array, batch_size=16, language=language, chunk_size=chunk_size)
-    return result
-
-
-def _lookup_model(model_id: str, device: str, compute_type: str) -> 'FasterWhisperPipeline':
-    import whisperx
-
-    key = (model_id, device, compute_type)
-    if key not in _model_cache:
-        model = whisperx.load_model(model_id, device, compute_type=compute_type)
-        _model_cache[key] = model
-    return _model_cache[key]
-
-
-_model_cache: dict[tuple[str, str, str], 'FasterWhisperPipeline'] = {}
-
-
-__all__ = local_public_names(__name__)
-
-
-def __dir__():
-    return __all__

pixeltable/ext/functions/yolox.py

@@ -1,157 +0,0 @@
-import logging
-from pathlib import Path
-from typing import TYPE_CHECKING, Iterable, Iterator
-from urllib.request import urlretrieve
-
-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]
-
-_logger = logging.getLogger('pixeltable')
-
-
-@pxt.udf(batch_size=4)
-def yolox(images: Batch[PIL.Image.Image], *, model_id: str, threshold: float = 0.5) -> Batch[dict]:
-    """
-    Computes YOLOX object detections for the specified image. `model_id` should reference one of the models
-    defined in the [YOLOX documentation](https://github.com/Megvii-BaseDetection/YOLOX).
-
-    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
-
-    __Requirements__:
-
-    - `pip install git+https://github.com/Megvii-BaseDetection/YOLOX`
-
-    Args:
-        model_id: one of: `yolox_nano`, `yolox_tiny`, `yolox_s`, `yolox_m`, `yolox_l`, `yolox_x`
-        threshold: the threshold for object detection
-
-    Returns:
-        A dictionary containing the output of the object detection model.
-
-    Examples:
-        Add a computed column that applies the model `yolox_m` to an existing
-        Pixeltable column `tbl.image` of the table `tbl`:
-
-        >>> tbl['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)
-
-    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
-
-
-@pxt.udf
-def yolo_to_coco(detections: dict) -> list:
-    """
-    Converts the output of a YOLOX object detection model to COCO format.
-
-    YOLOX is part of the `pixeltable.ext` package: long-term support in Pixeltable is not guaranteed.
-
-    Args:
-        detections: The output of a YOLOX object detection model, as returned by `yolox`.
-
-    Returns:
-        A dictionary containing the data from `detections`, converted to COCO format.
-
-    Examples:
-        Add a computed column that converts the output `tbl.detections` to COCO format, where `tbl.image`
-        is the image for which detections were computed:
-
-        >>> tbl['detections'] = yolox(tbl.image, model_id='yolox_m', threshold=0.8)
-        ... tbl['detections_coco'] = yolo_to_coco(tbl.detections)
-    """
-    bboxes, labels = detections['bboxes'], detections['labels']
-    num_annotations = len(detections['bboxes'])
-    assert num_annotations == len(detections['labels'])
-    result = []
-    for i in range(num_annotations):
-        bbox = bboxes[i]
-        ann = {
-            'bbox': [round(bbox[0]), round(bbox[1]), round(bbox[2] - bbox[0]), round(bbox[3] - bbox[1])],
-            'category': labels[i],
-        }
-        result.append(ann)
-    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:
-        image = normalize_image_mode(image)
-        image_transform, _ = _val_transform(np.array(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
-
-    key = (model_id, device)
-    if key in _model_cache:
-        return _model_cache[key]
-
-    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)
-
-    exp = get_exp(exp_name=model_id)
-    model = exp.get_model().to(device)
-
-    model.eval()
-    model.head.training = False
-    model.training = False
-
-    # Load in the weights from training
-    weights = torch.load(weights_file, map_location=torch.device(device))
-    model.load_state_dict(weights['model'])
-
-    _model_cache[key] = (model, exp)
-    return model, exp
-
-
-_model_cache: dict[tuple[str, str], tuple['YOLOX', 'Exp']] = {}
-
-
-__all__ = local_public_names(__name__)
-
-
-def __dir__():
-    return __all__