pixeltable 0.1.2__py3-none-any.whl → 0.2.1__py3-none-any.whl

pixeltable/utils/__init__.py

@@ -1,46 +1,17 @@
-from typing import Optional
-from pathlib import Path
-import glob
-
-from pixeltable.env import Env
-
-
-def is_computed_img_path(path: str) -> bool:
-    try:
-        _ = Path(path).relative_to(Env.get().img_dir)
-        return True
-    except ValueError:
-        return False
-
-def get_computed_img_path(tbl_id: int, col_id: int, version: int, rowid: int) -> str:
-    return Env.get().img_dir / f'img_{tbl_id}_{col_id}_{version}_{rowid}.jpg'
-
-def get_extracted_frame_path(tbl_id: int, video_col_id: int, version: int, offset: int) -> str:
-    return Env.get().img_dir / f'frame_{tbl_id}_{video_col_id}_{version}_{offset}'
-
-def computed_imgs(
-        tbl_id: Optional[int] = None, col_id: Optional[int] = None, version: Optional[int] = None) -> int:
-    path = f'{Env.get().img_dir}/img_'
-    path += f'{tbl_id}_' if tbl_id is not None else '*_'
-    path += f'{col_id}_' if col_id is not None else '*_'
-    path += f'{version}_' if version is not None else '*_'
-    path += '*'
-    names = glob.glob(path)
-    return names
-
-def computed_img_count(
-        tbl_id: Optional[int] = None, col_id: Optional[int] = None, version: Optional[int] = None) -> int:
-    return len(computed_imgs(tbl_id=tbl_id, col_id=col_id, version=version))
-
-def extracted_frames(tbl_id: Optional[int] = None, version: Optional[int] = None) -> int:
-    path = f'{Env.get().img_dir}/frame_'
-    path += f'{tbl_id}_' if tbl_id is not None else '*_'
-    path += '*_'  # video_col_id
-    path += f'{version}_' if version is not None else '*_'
-    path += '*_'  # offset
-    path += '*'  # running frame index
-    names = glob.glob(path)
-    return names
-
-def extracted_frame_count(tbl_id: Optional[int] = None, version: Optional[int] = None) -> int:
-    return len(extracted_frames(tbl_id, version))
+def print_perf_counter_delta(delta: float) -> str:
+    """Prints a performance counter delta in a human-readable format.
+
+    Args:
+        delta: delta in seconds
+
+    Returns:
+        Human-readable string
+    """
+    if delta < 1e-6:
+        return f'{delta * 1e9:.2f} ns'
+    elif delta < 1e-3:
+        return f'{delta * 1e6:.2f} us'
+    elif delta < 1:
+        return f'{delta * 1e3:.2f} ms'
+    else:
+        return f'{delta:.2f} s'

pixeltable/utils/clip.py

@@ -1,21 +1,18 @@
 import numpy as np
 import PIL.Image
-import clip
-import torch
-import PIL.Image
 
+import pixeltable.func as func
+from pixeltable.env import Env
 
-_device = 'cuda' if torch.cuda.is_available() else 'cpu'
-_model, _preprocess = clip.load("ViT-B/32", device=_device)
 
-def encode_image(img: PIL.Image.Image) -> np.ndarray:
-    preprocessed = _preprocess(img).unsqueeze(0).to(_device)
-    features = _model.encode_image(preprocessed)
-    val = features.numpy(force=True).squeeze()
-    return val
+def embed_image(img: PIL.Image.Image) -> np.ndarray:
+    from pixeltable.functions.nos.image_embedding import openai_clip
+    model_info = openai_clip.model_spec
+    result = Env.get().nos_client.Run(task=model_info.task, model_name=model_info.name, images=[img.resize((224, 224))])
+    return result['embedding'].squeeze(0)
 
-def encode_text(txt: str) -> np.ndarray:
-    preprocessed = clip.tokenize([txt]).to(_device)
-    features = _model.encode_text(preprocessed)
-    val = features.numpy(force=True).squeeze()
-    return val
+def embed_text(text: str) -> np.ndarray:
+    from pixeltable.functions.nos.text_embedding import openai_clip
+    model_info = openai_clip.model_spec
+    result = Env.get().nos_client.Run(task=model_info.task, model_name=model_info.name, texts=[text])
+    return result['embedding'].squeeze(0)

pixeltable/utils/coco.py

@@ -0,0 +1,136 @@
+from typing import List, Dict, Any, Set
+from pathlib import Path
+import json
+
+import PIL
+
+import pixeltable.exceptions as excs
+
+
+format_msg = """
+
+Required format:
+{
+    'image': PIL.Image.Image,
+    'annotations': [
+        {
+            'bbox': [x: int, y: int, w: int, h: int],
+            'category': str | int,
+        },
+        ...
+    ],
+}
+"""
+
+def _verify_input_dict(input_dict: Dict[str, Any]) -> None:
+    """Verify that input_dict is a valid input dict for write_coco_dataset()"""
+    if not isinstance(input_dict, dict):
+        raise excs.Error(f'Expected dict, got {input_dict}{format_msg}')
+    if 'image' not in input_dict:
+        raise excs.Error(f'Missing key "image" in input dict: {input_dict}{format_msg}')
+    if not isinstance(input_dict['image'], PIL.Image.Image):
+        raise excs.Error(f'Value for "image" is not a PIL.Image.Image: {input_dict}{format_msg}')
+    if  'annotations' not in input_dict:
+        raise excs.Error(f'Missing key "annotations" in input dict: {input_dict}{format_msg}')
+    if not isinstance(input_dict['annotations'], list):
+        raise excs.Error(f'Value for "annotations" is not a list: {input_dict}{format_msg}')
+    for annotation in input_dict['annotations']:
+        if not isinstance(annotation, dict):
+            raise excs.Error(f'Annotation is not a dict: {annotation}{format_msg}')
+        if 'bbox' not in annotation:
+            raise excs.Error(f'Missing key "bbox" in annotation: {annotation}{format_msg}')
+        if not isinstance(annotation['bbox'], list):
+            raise excs.Error(f'Value for "bbox" is not a list [x, y, w, h]: {annotation}{format_msg}')
+        if len(annotation['bbox']) != 4 or not all(isinstance(x, int) for x in annotation['bbox']):
+            raise excs.Error(f'Key "bbox" is not a list [x, y, w, h] of ints: {annotation}{format_msg}')
+        if 'category' not in annotation:
+            raise excs.Error(f'Missing key "category" in annotation: {annotation}{format_msg}')
+        if not isinstance(annotation['category'], (str, int)):
+            raise excs.Error(f'Value for "category" is not a str or int: {annotation}{format_msg}')
+
+def write_coco_dataset(df: 'pixeltable.DataFrame', dest_path: Path) -> Path:
+    """Export a DataFrame result set as a COCO dataset in dest_path and return the path of the data.json file."""
+    # TODO: validate schema
+    if len(df._select_list_exprs) != 1 or not df._select_list_exprs[0].col_type.is_json_type():
+        raise excs.Error(f'Expected exactly one json-typed column in select list: {df._select_list_exprs}')
+    input_dict_slot_idx = -1  # df._select_list_exprs[0].slot_idx isn't valid until _exec()
+
+    # create output dir
+    assert not dest_path.exists()
+    dest_path.mkdir(parents=False)
+    images_dir = dest_path / 'images'
+    images_dir.mkdir()
+
+    images: List[Dict[str, Any]] = []
+    img_id = -1
+    annotations: List[Dict[str, Any]] = []
+    ann_id = -1
+    categories: Set[Any] = set()
+    for input_row in df._exec():
+        if input_dict_slot_idx == -1:
+            input_dict_expr = df._select_list_exprs[0]
+            input_dict_slot_idx = input_dict_expr.slot_idx
+            input_dict = input_row[input_dict_slot_idx]
+            _verify_input_dict(input_dict)
+
+            # we want to know the slot idx of the image used in the input dict, so that we can check whether we
+            # already have a local path for it
+            input_dict_dependencies = input_dict_expr.dependencies()
+            img_slot_idx = next((e.slot_idx for e in input_dict_dependencies if e.col_type.is_image_type()), None)
+            assert img_slot_idx is not None
+        else:
+            input_dict = input_row[input_dict_slot_idx]
+            _verify_input_dict(input_dict)
+
+        # create image record
+        img_id += 1
+
+        # get a local path for the image
+        img = input_dict['image']
+        if input_row.file_paths[img_slot_idx] is not None:
+            # we already have a local path
+            img_path = Path(input_row.file_paths[img_slot_idx])
+            # TODO: if the path leads to our tmp dir, we need to move the file
+        else:
+            # we need to create a local path
+            img_path = images_dir / f'{img_id}.jpg'
+            img.save(img_path)
+
+        images.append({
+            'id': img_id,
+            'file_name': str(img_path),
+            'width': img.width,
+            'height': img.height,
+        })
+
+        # create annotation records for this image
+        for annotation in input_dict['annotations']:
+            ann_id += 1
+            x, y, w, h = annotation['bbox']
+            category = annotation['category']
+            categories.add(category)
+            annotations.append({
+                'id': ann_id,
+                'image_id': img_id,
+                # we use the category name here and fix it up at the end, when we have assigned category ids
+                'category_id': category,
+                'bbox': annotation['bbox'],
+                'area': w * h,
+                'iscrowd': 0,
+            })
+
+    # replace category names with ids
+    category_ids = {category: id for id, category in enumerate(sorted(list(categories)))}
+    for annotation in annotations:
+        annotation['category_id'] = category_ids[annotation['category_id']]
+
+    result = {
+        'images': images,
+        'annotations': annotations,
+        'categories': [{'id': id, 'name': category} for category, id in category_ids.items()],
+    }
+    output_path = dest_path / 'data.json'
+    with open(output_path, 'w') as f:
+        json.dump(result, f)
+    return output_path
+

pixeltable/utils/documents.py

@@ -0,0 +1,39 @@
+from typing import Optional, Dict
+import dataclasses
+
+import pixeltable.type_system as ts
+
+
+@dataclasses.dataclass
+class DocumentHandle:
+    format: ts.DocumentType.DocumentFormat
+    bs_doc: Optional['bs4.BeautifulSoup'] = None
+    md_ast: Optional[Dict] = None
+
+
+def get_document_handle(s: str) -> Optional[DocumentHandle]:
+    bs_doc = get_html_handle(s)
+    if bs_doc is not None:
+        return DocumentHandle(format=ts.DocumentType.DocumentFormat.HTML, bs_doc=bs_doc)
+    md_ast = get_markdown_handle(s)
+    if md_ast is not None:
+        return DocumentHandle(format=ts.DocumentType.DocumentFormat.MD, md_ast=md_ast)
+    return None
+
+def get_html_handle(s: str) -> Optional['bs4.BeautifulSoup']:
+    import bs4
+    try:
+        doc = bs4.BeautifulSoup(s, 'html.parser')
+    except Exception as e:
+        return None
+    if doc.find() is None:
+        return None
+    return doc
+
+def get_markdown_handle(s: str) -> Optional[Dict]:
+    import mistune
+    try:
+        md_ast = mistune.create_markdown(renderer=None)
+        return md_ast(s)
+    except Exception as e:
+        return None

pixeltable/utils/filecache.py

@@ -0,0 +1,195 @@
+from __future__ import annotations
+from typing import Optional, List, Tuple, Dict
+from collections import OrderedDict, defaultdict, namedtuple
+import os
+import glob
+from pathlib import Path
+from time import time
+import logging
+from uuid import UUID
+import hashlib
+
+from pixeltable.env import Env
+
+
+_logger = logging.getLogger('pixeltable')
+
+class CacheEntry:
+    def __init__(self, key: str, tbl_id: UUID, col_id: int, size: int, last_accessed_ts: int, ext: str):
+        self.key = key
+        self.tbl_id = tbl_id
+        self.col_id = col_id
+        self.size = size
+        self.last_accessed_ts = last_accessed_ts
+        self.ext = ext
+
+    def path(self) -> Path:
+        return Env.get().file_cache_dir / f'{self.tbl_id.hex}_{self.col_id}_{self.key}{self.ext}'
+
+    @classmethod
+    def from_file(cls, path: Path) -> CacheEntry:
+        components = path.stem.split('_')
+        assert len(components) == 3
+        tbl_id = UUID(components[0])
+        col_id = int(components[1])
+        key = components[2]
+        file_info = os.stat(str(path))
+        return cls(key, tbl_id, col_id, file_info.st_size, file_info.st_mtime, path.suffix)
+
+
+class FileCache:
+    """
+    A local cache of external (eg, S3) file references in cells of a stored table (ie, table or view).
+
+    Cache entries are identified by a hash of the file url and stored in Env.filecache_dir. The time of last
+    access of a cache entries is its file's mtime.
+
+    TODO:
+    - enforce a maximum capacity with LRU eviction
+    - implement MRU eviction for queries that exceed the capacity
+    """
+    _instance: Optional[FileCache] = None
+    ColumnStats = namedtuple('FileCacheColumnStats', ['tbl_id', 'col_id', 'num_files', 'total_size'])
+    CacheStats = namedtuple(
+        'FileCacheStats', ['total_size', 'num_requests', 'num_hits', 'num_evictions', 'column_stats'])
+
+    @classmethod
+    def get(cls) -> FileCache:
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    def __init__(self):
+        self.cache: OrderedDict[str, CacheEntry] = OrderedDict()  # ordered by entry.last_accessed_ts
+        self.total_size = 0
+        #self.capacity = Env.get().max_filecache_size
+        self.num_requests = 0
+        self.num_hits = 0
+        self.num_evictions = 0
+        paths = glob.glob(str(Env.get().file_cache_dir / '*'))
+        entries = [CacheEntry.from_file(Path(path_str)) for path_str in paths]
+        # we need to insert entries in order of last_accessed_ts
+        entries.sort(key=lambda e: e.last_accessed_ts)
+        for entry in entries:
+            self.cache[entry.key] = entry
+            self.total_size += entry.size
+
+    def avg_file_size(self) -> int:
+        if len(self.cache) == 0:
+            return 0
+        return int(self.total_size / len(self.cache))
+
+    def num_files(self, tbl_id: Optional[UUID] = None) -> int:
+        if tbl_id is None:
+            return len(self.cache)
+        entries = [e for e in self.cache.values() if e.tbl_id == tbl_id]
+        return len(entries)
+
+    def clear(self, tbl_id: Optional[UUID] = None, capacity: Optional[int] = None) -> None:
+        """
+        For testing purposes: allow resetting capacity and stats.
+        """
+        self.num_requests, self.num_hits, self.num_evictions = 0, 0, 0
+        entries = list(self.cache.values())  # list(): avoid dealing with values() return type
+        if tbl_id is not None:
+            entries = [e for e in entries if e.tbl_id == tbl_id]
+            _logger.debug(f'clearing {len(entries)} entries from file cache for table {tbl_id}')
+        else:
+            _logger.debug(f'clearing {len(entries)} entries from file cache')
+        for entry in entries:
+            del self.cache[entry.key]
+            self.total_size -= entry.size
+            os.remove(entry.path())
+        # if capacity is not None:
+        #     self.capacity = capacity
+        # else:
+        #     # need to reset to default
+        #     self.capacity = Env.get().max_filecache_size
+        # _logger.debug(f'setting file cache capacity to {self.capacity}')
+
+    def _url_hash(self, url: str) -> str:
+        h = hashlib.sha256()
+        h.update(url.encode())
+        return h.hexdigest()
+
+    def lookup(self, url: str) -> Optional[Path]:
+        self.num_requests += 1
+        key = self._url_hash(url)
+        entry = self.cache.get(key, None)
+        if entry is None:
+            _logger.debug(f'file cache miss for {url}')
+            return None
+        # update mtime and cache
+        path = entry.path()
+        path.touch(exist_ok=True)
+        file_info = os.stat(str(path))
+        entry.last_accessed_ts = file_info.st_mtime
+        self.cache.move_to_end(key, last=True)
+        self.num_hits += 1
+        _logger.debug(f'file cache hit for {url}')
+        return path
+
+    # def can_admit(self, query_ts: int) -> bool:
+    #     if self.total_size + self.avg_file_size <= self.capacity:
+    #         return True
+    #     assert len(self.cache) > 0
+    #     # check whether we can evict the current lru entry
+    #     lru_entry = next(iter(self.cache.values()))
+    #     if lru_entry.last_accessed_ts >= query_ts:
+    #         # the current query brought this entry in: we're not going to evict it
+    #         return False
+    #     return True
+
+    def add(self, tbl_id: UUID, col_id: int, url: str, path: Path) -> Path:
+        """Adds url at 'path' to cache and returns its new path.
+        'path' will not be accessible after this call. Retains the extension of 'path'.
+        """
+        file_info = os.stat(str(path))
+        _ = time()
+        #if self.total_size + file_info.st_size > self.capacity:
+        if False:
+            if len(self.cache) == 0:
+                # nothing to evict
+                return
+            # evict entries until we're below the limit or until we run into entries the current query brought in
+            while True:
+                lru_entry = next(iter(self.cache.values()))
+                if lru_entry.last_accessed_ts >= query_ts:
+                    # the current query brought this entry in: switch to MRU and ignore this put()
+                    _logger.debug('file cache switched to MRU')
+                    return
+                self.cache.popitem(last=False)
+                self.total_size -= lru_entry.size
+                self.num_evictions += 1
+                os.remove(str(lru_entry.path()))
+                _logger.debug(f'evicted entry for cell {lru_entry.cell_id} from file cache')
+                if self.total_size + file_info.st_size <= self.capacity:
+                    break
+
+        key = self._url_hash(url)
+        assert key not in self.cache
+        entry = CacheEntry(key, tbl_id, col_id, file_info.st_size, file_info.st_mtime, path.suffix)
+        self.cache[key] = entry
+        self.total_size += entry.size
+        new_path = entry.path()
+        os.rename(str(path), str(new_path))
+        _logger.debug(f'added entry for cell {url} to file cache')
+        return new_path
+
+    def stats(self) -> CacheStats:
+        # collect column stats
+        # (tbl_id, col_id) -> (num_files, total_size)
+        d: Dict[Tuple[int, int], List[int]] = defaultdict(lambda: [0, 0])
+        for entry in self.cache.values():
+            t = d[(entry.tbl_id, entry.col_id)]
+            t[0] += 1
+            t[1] += entry.size
+        col_stats = [
+            self.ColumnStats(tbl_id, col_id, num_files, size) for (tbl_id, col_id), (num_files, size) in d.items()
+        ]
+        col_stats.sort(key=lambda e: e[3], reverse=True)
+        return self.CacheStats(self.total_size, self.num_requests, self.num_hits, self.num_evictions, col_stats)
+
+    def debug_print(self) -> None:
+        for entry in self.cache.values():
+            print(f'CacheEntry: tbl_id={entry.tbl_id}, col_id={entry.col_id}, size={entry.size}')

pixeltable/utils/help.py

@@ -0,0 +1,11 @@
+from typing import Any
+
+import pixeltable.func as func
+
+
+def help(obj: Any) -> None:
+    """Returns help text for the given object."""
+    if isinstance(obj, func.Function):
+        print(obj.help_str())
+    else:
+        print(__builtins__.help(obj))

pixeltable/utils/media_store.py

@@ -0,0 +1,76 @@
+import glob
+import os
+import re
+import shutil
+import uuid
+from typing import Optional, List, Tuple, Dict
+from pathlib import Path
+from collections import defaultdict
+from uuid import UUID
+
+from pixeltable.env import Env
+
+
+class MediaStore:
+    """
+    Utilities to manage media files stored in Env.media_dir
+
+    Media file names are a composite of: table id, column id, version, uuid:
+    the table id/column id/version are redundant but useful for identifying all files for a table
+    or all files created for a particular version of a table
+    """
+    pattern = re.compile(r'([0-9a-fA-F]+)_(\d+)_(\d+)_([0-9a-fA-F]+)')  # tbl_id, col_id, version, uuid
+
+    @classmethod
+    def prepare_media_path(cls, tbl_id: UUID, col_id: int, version: int, ext: Optional[str] = None) -> Path:
+        """
+        Construct a new, unique Path name for a persisted media file, and create the parent directory
+        for the new Path if it does not already exist. The Path will reside in
+        the environment's media_dir.
+        """
+        id_hex = uuid.uuid4().hex
+        parent = Env.get().media_dir / tbl_id.hex / id_hex[0:2] / id_hex[0:4]
+        parent.mkdir(parents=True, exist_ok=True)
+        return parent / f'{tbl_id.hex}_{col_id}_{version}_{id_hex}{ext or ""}'
+
+    @classmethod
+    def delete(cls, tbl_id: UUID, version: Optional[int] = None) -> None:
+        """Delete all files belonging to tbl_id. If version is not None, delete
+        only those files belonging to the specified version."""
+        assert tbl_id is not None
+        if version is None:
+            # Remove the entire folder for this table id.
+            path = Env.get().media_dir / tbl_id.hex
+            if path.exists():
+                shutil.rmtree(path)
+        else:
+            # Remove only the elements for the specified version.
+            paths = glob.glob(str(Env.get().media_dir / tbl_id.hex) + f'/**/{tbl_id.hex}_*_{version}_*', recursive=True)
+            for path in paths:
+                os.remove(path)
+
+    @classmethod
+    def count(cls, tbl_id: UUID) -> int:
+        """
+        Return number of files for given tbl_id.
+        """
+        paths = glob.glob(str(Env.get().media_dir / tbl_id.hex) + f'/**/{tbl_id.hex}_*', recursive=True)
+        return len(paths)
+
+    @classmethod
+    def stats(cls) -> List[Tuple[int, int, int, int]]:
+        paths = glob.glob(str(Env.get().media_dir) + "/**", recursive=True)
+        # key: (tbl_id, col_id), value: (num_files, size)
+        d: Dict[Tuple[UUID, int], List[int]] = defaultdict(lambda: [0, 0])
+        for p in paths:
+            if not os.path.isdir(p):
+                matched = re.match(cls.pattern, Path(p).name)
+                assert matched is not None
+                tbl_id, col_id = UUID(hex=matched[1]), int(matched[2])
+                file_info = os.stat(p)
+                t = d[(tbl_id, col_id)]
+                t[0] += 1
+                t[1] += file_info.st_size
+        result = [(tbl_id, col_id, num_files, size) for (tbl_id, col_id), (num_files, size) in d.items()]
+        result.sort(key=lambda e: e[3], reverse=True)
+        return result