pixeltable 0.2.14__py3-none-any.whl → 0.2.16__py3-none-any.whl

pixeltable/functions/{eval.py → vision.py}

@@ -1,11 +1,29 @@
-from typing import List, Tuple, Dict
+"""
+Pixeltable [UDFs](https://pixeltable.readme.io/docs/user-defined-functions-udfs) for Computer Vision.
+
+Example:
+```python
+import pixeltable as pxt
+from pixeltable.functions import vision as pxtv
+
+t = pxt.get_table(...)
+t.select(pxtv.draw_bounding_boxes(t.img, boxes=t.boxes, label=t.labels)).collect()
+```
+"""
+
+import colorsys
+import hashlib
+import random
 from collections import defaultdict
-import sys
+from typing import Optional, Union, Any
 
+import PIL.Image
+import PIL.Image
 import numpy as np
 
-import pixeltable.type_system as ts
 import pixeltable.func as func
+import pixeltable.type_system as ts
+from pixeltable.utils.code import local_public_names
 
 
 # TODO: figure out a better submodule structure
@@ -14,7 +32,7 @@ import pixeltable.func as func
 # 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:
+def __calculate_bboxes_area(bboxes: np.ndarray) -> np.ndarray:
     """Calculate area of bounding boxes.
 
     Args:
@@ -31,7 +49,7 @@ def calculate_bboxes_area(bboxes: np.ndarray) -> np.ndarray:
 # 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:
+def __calculate_overlaps(bboxes1: np.ndarray, bboxes2: np.ndarray) -> np.ndarray:
     """Calculate the overlap between each bbox of bboxes1 and bboxes2.
 
     Args:
@@ -58,8 +76,8 @@ def calculate_overlaps(bboxes1: np.ndarray, bboxes2: np.ndarray) -> np.ndarray:
         exchange = False
 
     # Calculate the bboxes area.
-    area1 = calculate_bboxes_area(bboxes1)
-    area2 = calculate_bboxes_area(bboxes2)
+    area1 = __calculate_bboxes_area(bboxes1)
+    area2 = __calculate_bboxes_area(bboxes2)
     eps = np.finfo(np.float32).eps
 
     for i in range(bboxes1.shape[0]):
@@ -80,9 +98,9 @@ def calculate_overlaps(bboxes1: np.ndarray, bboxes2: np.ndarray) -> np.ndarray:
 # 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(
+def __calculate_image_tpfp(
     pred_bboxes: np.ndarray, pred_scores: np.ndarray, gt_bboxes: np.ndarray, min_iou: float
-) -> Tuple[np.ndarray, np.ndarray]:
+) -> tuple[np.ndarray, np.ndarray]:
     """Calculate the true positive and false positive on an image.
 
     Args:
@@ -95,11 +113,8 @@ def calculate_image_tpfp(
 
     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.
+        tp: Shape (N,), the true positive flag of each predicted bbox on this image.
+        fp: 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`.
@@ -121,7 +136,7 @@ def calculate_image_tpfp(
 
     # Step 4. Calculate the IoUs between the predicted bboxes and the
     # ground truth bboxes.
-    ious = calculate_overlaps(pred_bboxes, gt_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.
@@ -160,14 +175,17 @@ def calculate_image_tpfp(
     ],
 )
 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:
+    pred_bboxes: list[list[int]],
+    pred_labels: list[int],
+    pred_scores: list[float],
+    gt_bboxes: list[list[int]],
+    gt_labels: list[int],
+) -> dict:
+    """
+    Evaluates the performance of a set of predicted bounding boxes against a set of ground truth bounding boxes.
+    """
     class_idxs = list(set(pred_labels + gt_labels))
-    result: List[Dict] = []
+    result: list[dict] = []
     pred_bboxes_arr = np.asarray(pred_bboxes)
     pred_classes_arr = np.asarray(pred_labels)
     pred_scores_arr = np.asarray(pred_scores)
@@ -177,7 +195,7 @@ 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(
             {
@@ -194,17 +212,21 @@ def eval_detections(
 
 @func.uda(update_types=[ts.JsonType()], value_type=ts.JsonType(), allows_std_agg=True, allows_window=False)
 class mean_ap(func.Aggregator):
+    """
+    Calculates the mean average precision (mAP) over
+    [`eval_detections()`][pixeltable.functions.vision.eval_detections] results.
+    """
     def __init__(self):
-        self.class_tpfp: Dict[int, List[Dict]] = defaultdict(list)
+        self.class_tpfp: dict[int, list[dict]] = defaultdict(list)
 
-    def update(self, eval_dicts: List[Dict]) -> None:
+    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:
+    def value(self) -> dict:
         eps = np.finfo(np.float32).eps
-        result: Dict[int, float] = {}
+        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)
@@ -225,3 +247,120 @@ class mean_ap(func.Aggregator):
             ap = np.sum((mrec[ind + 1] - mrec[ind]) * mpre[ind + 1])
             result[class_idx] = ap.item()
         return result
+
+
+def _create_label_colors(labels: list[Any]) -> dict[Any, str]:
+    """
+    Create random colors for labels such that a particular label always gets the same color.
+
+    Returns:
+        dict mapping labels to colors
+    """
+    distinct_labels = set(labels)
+    result: dict[Any, str] = {}
+    for label in distinct_labels:
+        # consistent hash for the label
+        label_hash = int(hashlib.md5(str(label).encode()).hexdigest(), 16)
+        hue = (label_hash % 360) / 360.0
+        rgb = colorsys.hsv_to_rgb(hue, 0.7, 0.95)
+        hex_color = '#{:02x}{:02x}{:02x}'.format(int(rgb[0]*255), int(rgb[1]*255), int(rgb[2]*255))
+        result[label] = hex_color
+    return result
+
+
+@func.udf
+def draw_bounding_boxes(
+        img: PIL.Image.Image,
+        boxes: list[list[int]],
+        labels: Optional[list[Any]] = None,
+        color: Optional[str] = None,
+        box_colors: Optional[list[str]] = None,
+        fill: bool  = False,
+        width: int = 1,
+        font: Optional[str] = None,
+        font_size: Optional[int] = None,
+) -> PIL.Image.Image:
+    """
+    Draws bounding boxes on the given image.
+
+    Labels can be any type that supports `str()` and is hashable (e.g., strings, ints, etc.).
+
+    Colors can be specified as common HTML color names (e.g., 'red') supported by PIL's
+    [`ImageColor`](https://pillow.readthedocs.io/en/stable/reference/ImageColor.html#imagecolor-module) module or as
+    RGB hex codes (e.g., '#FF0000').
+
+    If no colors are specified, this function randomly assigns each label a specific color based on a hash of the label.
+
+    Args:
+        img: The image on which to draw the bounding boxes.
+        boxes: List of bounding boxes, each represented as [xmin, ymin, xmax, ymax].
+        labels: List of labels for each bounding box.
+        color: Single color to be used for all bounding boxes and labels.
+        box_colors: List of colors, one per bounding box.
+        fill: Whether to fill the bounding boxes with color.
+        width: Width of the bounding box borders.
+        font: Name of a system font or path to a TrueType font file, as required by
+            [`PIL.ImageFont.truetype()`](https://pillow.readthedocs.io/en/stable/reference/ImageFont.html#PIL.ImageFont.truetype).
+            If `None`, uses the default provided by
+            [`PIL.ImageFont.load_default()`](https://pillow.readthedocs.io/en/stable/reference/ImageFont.html#PIL.ImageFont.load_default).
+        font_size: Size of the font used for labels in points. Only used in conjunction with non-`None` `font` argument.
+
+    Returns:
+        The image with bounding boxes drawn on it.
+    """
+    color_params = sum([color is not None, box_colors is not None])
+    if color_params > 1:
+        raise ValueError("Only one of 'color' or 'box_colors' can be set")
+
+    # ensure the number of labels matches the number of boxes
+    num_boxes = len(boxes)
+    if labels is None:
+        labels = [None] * num_boxes
+    elif len(labels) != num_boxes:
+        raise ValueError('Number of boxes and labels must match')
+
+    DEFAULT_COLOR = 'white'
+    if box_colors is not None:
+        if len(box_colors) != num_boxes:
+            raise ValueError('Number of boxes and box colors must match')
+    else:
+        if color is not None:
+            box_colors = [color] * num_boxes
+        else:
+            label_colors = _create_label_colors(labels)
+            box_colors = [label_colors[label] for label in labels]
+
+    from PIL import ImageDraw, ImageFont, ImageColor
+    # set default font if not provided
+    if font is None:
+        txt_font = ImageFont.load_default()
+    else:
+        txt_font = ImageFont.truetype(font=font, size=font_size or 10)
+
+    img_to_draw = img.copy()
+    draw = ImageDraw.Draw(img_to_draw, 'RGBA' if fill else 'RGB')
+
+    for i, (bbox, label) in enumerate(zip(boxes, labels)):
+        # determine color for the current box and label
+        color = box_colors[i % len(box_colors)]
+
+        if fill:
+            rgb_color = ImageColor.getrgb(color)
+            fill_color = rgb_color + (100,)  # semi-transparent
+            draw.rectangle(bbox, outline=color, width=width, fill=fill_color)
+        else:
+            draw.rectangle(bbox, outline=color, width=width)
+
+        if label is not None:
+            label_str = str(label)
+            margin = width + 1
+            draw.text((bbox[0] + margin, bbox[1] + margin), label_str, fill=color, font=txt_font)
+
+    return img_to_draw
+
+
+__all__ = local_public_names(__name__)
+
+
+def __dir__():
+    return __all__

pixeltable/globals.py

@@ -1,15 +1,18 @@
 import dataclasses
 import logging
 from typing import Any, Optional, Union
+from uuid import UUID
 
 import pandas as pd
 import sqlalchemy as sql
+from pandas.io.formats.style import Styler
 from sqlalchemy.util.preloaded import orm
 
 import pixeltable.exceptions as excs
 import pixeltable.exprs as exprs
-from pixeltable import catalog, func, DataFrame
+from pixeltable import DataFrame, catalog, func
 from pixeltable.catalog import Catalog
+from pixeltable.dataframe import DataFrameResultSet
 from pixeltable.env import Env
 from pixeltable.iterators import ComponentIterator
 from pixeltable.metadata import schema
@@ -24,21 +27,25 @@ def init() -> None:
 
 def create_table(
     path_str: str,
-    schema: dict[str, Any],
+    schema_or_df: Union[dict[str, Any], DataFrame],
     *,
     primary_key: Optional[Union[str, list[str]]] = None,
     num_retained_versions: int = 10,
     comment: str = '',
-) -> catalog.InsertableTable:
-    """Create a new `InsertableTable`.
+) -> catalog.Table:
+    """Create a new base table.
 
     Args:
         path_str: Path to the table.
-        schema: dictionary mapping column names to column types, value expressions, or to column specifications.
+        schema_or_df: Either a dictionary that maps column names to column types, or a
+            [`DataFrame`][pixeltable.DataFrame] whose contents and schema will be used to pre-populate the table.
+        primary_key: An optional column name or list of column names to use as the primary key(s) of the
+            table.
         num_retained_versions: Number of versions of the table to retain.
+        comment: An optional comment; its meaning is user-defined.
 
     Returns:
-        The newly created table.
+        A handle to the newly created [`Table`][pixeltable.Table].
 
     Raises:
         Error: if the path already exists or is invalid.
@@ -46,12 +53,27 @@ def create_table(
     Examples:
         Create a table with an int and a string column:
 
-        >>> table = cl.create_table('my_table', schema={'col1': IntType(), 'col2': StringType()})
+        >>> table = pxt.create_table('my_table', schema={'col1': IntType(), 'col2': StringType()})
+
+        Create a table from a select statement over an existing table `tbl`:
+
+        >>> table = pxt.create_table('my_table', tbl.where(tbl.col1 < 10).select(tbl.col2))
     """
     path = catalog.Path(path_str)
     Catalog.get().paths.check_is_valid(path, expected=None)
     dir = Catalog.get().paths[path.parent]
 
+    df: Optional[DataFrame] = None
+    if isinstance(schema_or_df, dict):
+        schema = schema_or_df
+    elif isinstance(schema_or_df, DataFrame):
+        df = schema_or_df
+        schema = df.schema
+    elif isinstance(schema_or_df, DataFrameResultSet):
+        raise excs.Error('`schema_or_df` must be either a schema dictionary or a Pixeltable DataFrame. (Is there an extraneous call to `collect()`?)')
+    else:
+        raise excs.Error('`schema_or_df` must be either a schema dictionary or a Pixeltable DataFrame.')
+
     if len(schema) == 0:
         raise excs.Error(f'Table schema is empty: `{path_str}`')
 
@@ -63,15 +85,17 @@ def create_table(
         if not isinstance(primary_key, list) or not all(isinstance(pk, str) for pk in primary_key):
             raise excs.Error('primary_key must be a single column name or a list of column names')
 
-    tbl = catalog.InsertableTable.create(
+    tbl = catalog.InsertableTable._create(
         dir._id,
         path.name,
         schema,
+        df,
         primary_key=primary_key,
         num_retained_versions=num_retained_versions,
         comment=comment,
     )
     Catalog.get().paths[path] = tbl
+
     _logger.info(f'Created table `{path_str}`.')
     return tbl
 
@@ -87,25 +111,28 @@ def create_view(
     num_retained_versions: int = 10,
     comment: str = '',
     ignore_errors: bool = False,
-) -> catalog.View:
-    """Create a new `View`.
+) -> Optional[catalog.Table]:
+    """Create a view of an existing table object (which itself can be a view or a snapshot or a base table).
 
     Args:
         path_str: Path to the view.
-        base: Table (i.e., table or view or snapshot) or DataFrame to base the view on.
+        base: [`Table`][pixeltable.Table] (i.e., table or view or snapshot) or [`DataFrame`][pixeltable.DataFrame] to
+            base the view on.
         schema: dictionary mapping column names to column types, value expressions, or to column specifications.
         filter: predicate to filter rows of the base table.
         is_snapshot: Whether the view is a snapshot.
         iterator: The iterator to use for this view. If specified, then this view will be a one-to-many view of
             the base table.
         num_retained_versions: Number of versions of the view to retain.
+        comment: Optional comment for the view.
         ignore_errors: if True, fail silently if the path already exists or is invalid.
 
     Returns:
-        The newly created view.
+        A handle to the [`Table`][pixeltable.Table] representing the newly created view. If the path already
+        exists or is invalid and `ignore_errors=True`, returns `None`.
 
     Raises:
-        Error: if the path already exists or is invalid.
+        Error: if the path already exists or is invalid and `ignore_errors=False`.
 
     Examples:
         Create a view with an additional int and a string column and a filter:
@@ -140,7 +167,7 @@ def create_view(
         Catalog.get().paths.check_is_valid(path, expected=None)
     except Exception as e:
         if ignore_errors:
-            return
+            return None
         else:
             raise e
     dir = Catalog.get().paths[path.parent]
@@ -152,7 +179,7 @@ def create_view(
     else:
         iterator_class, iterator_args = iterator
 
-    view = catalog.View.create(
+    view = catalog.View._create(
         dir._id,
         path.name,
         base=tbl_version_path,
@@ -170,16 +197,16 @@ def create_view(
 
 
 def get_table(path: str) -> catalog.Table:
-    """Get a handle to a table (including views and snapshots).
+    """Get a handle to an existing table or view or snapshot.
 
     Args:
         path: Path to the table.
 
     Returns:
-        A `InsertableTable` or `View` object.
+        A handle to the [`Table`][pixeltable.Table].
 
     Raises:
-        Error: If the path does not exist or does not designate a table.
+        Error: If the path does not exist or does not designate a table object.
 
     Examples:
         Get handle for a table in the top-level directory:
@@ -197,6 +224,7 @@ def get_table(path: str) -> catalog.Table:
     p = catalog.Path(path)
     Catalog.get().paths.check_is_valid(p, expected=catalog.Table)
     obj = Catalog.get().paths[p]
+    assert isinstance(obj, catalog.Table)
     return obj
 
 
@@ -230,15 +258,15 @@ def move(path: str, new_path: str) -> None:
 
 
 def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> None:
-    """Drop a table.
+    """Drop a table or view or snapshot.
 
     Args:
-        path: Path to the table.
+        path: Path to the [`Table`][pixeltable.Table].
         force: If `True`, will also drop all views or sub-views of this table.
         ignore_errors: Whether to ignore errors if the table does not exist.
 
     Raises:
-        Error: If the path does not exist or does not designate a table and ignore_errors is False.
+        Error: If the path does not exist or does not designate a table object and ignore_errors is False.
 
     Examples:
         >>> cl.drop_table('my_table')
@@ -256,7 +284,7 @@ def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> N
     tbl = cat.paths[path_obj]
     assert isinstance(tbl, catalog.Table)
     if len(cat.tbl_dependents[tbl._id]) > 0:
-        dependent_paths = [dep.path for dep in cat.tbl_dependents[tbl._id]]
+        dependent_paths = [dep._path for dep in cat.tbl_dependents[tbl._id]]
         if force:
             for dependent_path in dependent_paths:
                 drop_table(dependent_path, force=True)
@@ -268,14 +296,14 @@ def drop_table(path: str, force: bool = False, ignore_errors: bool = False) -> N
 
 
 def list_tables(dir_path: str = '', recursive: bool = True) -> list[str]:
-    """List the tables in a directory.
+    """List the [`Table`][pixeltable.Table]s in a directory.
 
     Args:
         dir_path: Path to the directory. Defaults to the root directory.
         recursive: Whether to list tables in subdirectories as well.
 
     Returns:
-        A list of table paths.
+        A list of [`Table`][pixeltable.Table] paths.
 
     Raises:
         Error: If the path does not exist or does not designate a directory.
@@ -297,7 +325,7 @@ def list_tables(dir_path: str = '', recursive: bool = True) -> list[str]:
     return [str(p) for p in Catalog.get().paths.get_children(path, child_type=catalog.Table, recursive=recursive)]
 
 
-def create_dir(path_str: str, ignore_errors: bool = False) -> catalog.Dir:
+def create_dir(path_str: str, ignore_errors: bool = False) -> Optional[catalog.Dir]:
     """Create a directory.
 
     Args:
@@ -325,6 +353,7 @@ def create_dir(path_str: str, ignore_errors: bool = False) -> catalog.Dir:
             session.add(dir_record)
             session.flush()
             assert dir_record.id is not None
+            assert isinstance(dir_record.id, UUID)
             dir = catalog.Dir(dir_record.id, parent._id, path.name)
             Catalog.get().paths[path] = dir
             session.commit()
@@ -333,7 +362,7 @@ def create_dir(path_str: str, ignore_errors: bool = False) -> catalog.Dir:
             return dir
     except excs.Error as e:
         if ignore_errors:
-            return
+            return None
         else:
             raise e
 
@@ -415,7 +444,7 @@ def list_dirs(path_str: str = '', recursive: bool = True) -> list[str]:
     return [str(p) for p in Catalog.get().paths.get_children(path, child_type=catalog.Dir, recursive=recursive)]
 
 
-def list_functions() -> pd.DataFrame:
+def list_functions() -> Styler:
     """Returns information about all registered functions.
 
     Returns:
@@ -436,7 +465,7 @@ def list_functions() -> pd.DataFrame:
             'Return Type': [str(f.signature.get_return_type()) for f in functions],
         }
     )
-    pd_df = pd_df.style.set_properties(**{'text-align': 'left'}).set_table_styles(
+    pd_df = pd_df.style.set_properties(None, **{'text-align': 'left'}).set_table_styles(
         [dict(selector='th', props=[('text-align', 'center')])]
     )  # center-align headings
     return pd_df.hide(axis='index')

pixeltable/io/external_store.py

@@ -217,17 +217,17 @@ class Project(ExternalStore, abc.ABC):
         resolved_col_mapping: dict[Column, str] = {}
 
         # Validate names
-        t_cols = table.column_names()
+        t_cols = set(table._schema.keys())
         for t_col, ext_col in col_mapping.items():
             if t_col not in t_cols:
                 if is_user_specified_col_mapping:
                     raise excs.Error(
-                        f'Column name `{t_col}` appears as a key in `col_mapping`, but Table `{table.name}` '
+                        f'Column name `{t_col}` appears as a key in `col_mapping`, but Table `{table._name}` '
                         'contains no such column.'
                     )
                 else:
                     raise excs.Error(
-                        f'Column `{t_col}` does not exist in Table `{table.name}`. Either add a column `{t_col}`, '
+                        f'Column `{t_col}` does not exist in Table `{table._name}`. Either add a column `{t_col}`, '
                         f'or specify a `col_mapping` to associate a different column with the external field `{ext_col}`.'
                     )
             if ext_col not in export_cols and ext_col not in import_cols:
@@ -238,13 +238,13 @@ class Project(ExternalStore, abc.ABC):
             col = table[t_col].col
             resolved_col_mapping[col] = ext_col
         # Validate column specs
-        t_col_types = table.column_types()
+        t_col_types = table._schema
         for t_col, ext_col in col_mapping.items():
             t_col_type = t_col_types[t_col]
             if ext_col in export_cols:
                 # Validate that the table column can be assigned to the external column
                 ext_col_type = export_cols[ext_col]
-                if not ext_col_type.is_supertype_of(t_col_type):
+                if not ext_col_type.is_supertype_of(t_col_type, ignore_nullable=True):
                     raise excs.Error(
                         f'Column `{t_col}` cannot be exported to external column `{ext_col}` (incompatible types; expecting `{ext_col_type}`)'
                     )
@@ -255,7 +255,7 @@ class Project(ExternalStore, abc.ABC):
                         f'Column `{t_col}` is a computed column, which cannot be populated from an external column'
                     )
                 ext_col_type = import_cols[ext_col]
-                if not t_col_type.is_supertype_of(ext_col_type):
+                if not t_col_type.is_supertype_of(ext_col_type, ignore_nullable=True):
                     raise excs.Error(
                         f'Column `{t_col}` cannot be imported from external column `{ext_col}` (incompatible types; expecting `{ext_col_type}`)'
                     )