pixeltable 0.1.0__py3-none-any.whl → 0.2.4__py3-none-any.whl

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

pixeltable/iterators/document.py

@@ -0,0 +1,311 @@
+from typing import Dict, Any, List, Tuple, Generator, Optional, Iterable
+import logging
+import dataclasses
+import enum
+
+from .base import ComponentIterator
+
+from pixeltable.type_system import ColumnType, DocumentType, StringType, IntType, JsonType
+from pixeltable.exceptions import Error
+from pixeltable.env import Env
+from pixeltable.utils.documents import get_document_handle
+
+
+_logger = logging.getLogger('pixeltable')
+
+
+class ChunkMetadata(enum.Enum):
+    TITLE = 1
+    HEADINGS = 2
+    SOURCELINE = 3
+
+
+class Separator(enum.Enum):
+    HEADING = 1
+    PARAGRAPH = 2
+    SENTENCE = 3
+    TOKEN_LIMIT = 4
+    CHAR_LIMIT = 5
+
+
+@dataclasses.dataclass
+class DocumentSectionMd:
+    """Metadata for a subsection of a document (ie, a structural element like a heading or paragraph)"""
+    source_line: int
+
+    # the stack of headings up to the most recently observed one;
+    # eg, if the most recent one was an h2, 'headings' would contain keys 1 and 2, but nothing below that
+    headings: Dict[int, str]
+
+
+@dataclasses.dataclass
+class DocumentSection:
+    """A single document chunk, according to some of the splitting criteria"""
+    text: Optional[str]
+    md: Optional[DocumentSectionMd]
+
+
+class DocumentSplitter(ComponentIterator):
+    """"Iterator over pieces of a document"""
+    MD_COLUMN_TYPES = {
+        ChunkMetadata.TITLE: StringType(),
+        ChunkMetadata.HEADINGS: JsonType(),
+        ChunkMetadata.SOURCELINE: IntType()
+    }
+
+    def __init__(
+            self, document: str, *, separators: str, limit: int = 0, overlap: int = 0, metadata: str = '',
+            html_skip_tags: List[str] = None, tiktoken_encoding: Optional[str] = 'cl100k_base',
+            tiktoken_target_model: Optional[str] = None
+    ):
+        import bs4
+        if html_skip_tags is None:
+            html_skip_tags = ['nav']
+        with open(document, 'r', encoding='utf8') as fh:
+            s = fh.read()
+            self._doc_handle = get_document_handle(s)
+            assert self._doc_handle is not None
+        self._separators = [Separator[s.upper()] for s in separators.split(',')]
+        self._md_fields = [ChunkMetadata[m.upper()] for m in metadata.split(',')] if len(metadata) > 0 else []
+        self._doc_title = \
+            self._doc_handle.bs_doc.title.get_text().strip() if self._doc_handle.bs_doc is not None else ''
+        self._limit = limit
+        self._skip_tags = html_skip_tags
+        self._overlap = overlap
+        self._tiktoken_encoding = tiktoken_encoding
+        self._tiktoken_target_model = tiktoken_target_model
+
+        # set up processing pipeline
+        if self._doc_handle.format == DocumentType.DocumentFormat.HTML:
+            assert self._doc_handle.bs_doc is not None
+            self._sections = self._html_sections()
+        else:
+            assert self._doc_handle.md_ast is not None
+            self._sections = self._markdown_sections()
+        if Separator.SENTENCE in self._separators:
+            self._sections = self._sentence_sections(self._sections)
+        if Separator.TOKEN_LIMIT in self._separators:
+            self._sections = self._token_chunks(self._sections)
+        if Separator.CHAR_LIMIT in self._separators:
+            self._sections = self._char_chunks(self._sections)
+
+    @classmethod
+    def input_schema(cls) -> Dict[str, ColumnType]:
+        return {
+            'document': DocumentType(nullable=False),
+            'separators': StringType(nullable=False),
+            'metadata': StringType(nullable=True),
+            'limit': IntType(nullable=True),
+            'overlap': IntType(nullable=True),
+            'skip_tags': StringType(nullable=True),
+            'tiktoken_encoding': StringType(nullable=True),
+            'tiktoken_target_model': StringType(nullable=True),
+        }
+
+    @classmethod
+    def output_schema(cls, *args: Any, **kwargs: Any) -> Tuple[Dict[str, ColumnType], List[str]]:
+        schema = {'text': StringType()}
+        if 'metadata' in kwargs and len(kwargs['metadata']) > 0:
+            md_fields = kwargs['metadata'].split(',')
+            for md_field in md_fields:
+                if not hasattr(ChunkMetadata, md_field.upper()):
+                    raise Error(f'Invalid metadata field {md_field}')
+                schema[md_field.lower()] = cls.MD_COLUMN_TYPES[ChunkMetadata[md_field.upper()]]
+
+        assert 'separators' in kwargs
+        separators = kwargs['separators'].split(',')
+        for separator in separators:
+            if not hasattr(Separator, separator.upper()):
+                raise Error(f'Invalid separator {separator}')
+
+        # check dependencies
+        if 'sentence' in separators:
+            Env.get().require_package('spacy')
+        if 'token_limit' in separators:
+            Env.get().require_package('tiktoken')
+
+        if 'limit' in kwargs or 'overlap' in kwargs:
+            if 'token_limit' not in separators and 'char_limit' not in separators:
+                raise Error('limit/overlap requires the "token_limit" or "char_limit" separator')
+            if 'limit' in kwargs and int(kwargs['limit']) <= 0:
+                raise Error('"limit" must be an integer > 0')
+            if 'overlap' in kwargs and int(kwargs['overlap']) < 0:
+                raise Error('"overlap" must be an integer >= 0')
+        if 'token_limit' in separators or 'char_limit' in separators:
+            if 'token_limit' in separators and 'char_limit' in separators:
+                raise Error('Cannot specify both "token_limit" and "char_limit" separators')
+            if 'limit' not in kwargs:
+                raise Error('limit is required with "token_limit"/"char_limit" separators')
+
+        return schema, []
+
+    def __next__(self) -> Dict[str, Any]:
+        while True:
+            section = next(self._sections)
+            if section.text is None:
+                continue
+            result = {'text': section.text}
+            for md_field in self._md_fields:
+                if md_field == ChunkMetadata.TITLE:
+                    result[md_field.name.lower()] = self._doc_title
+                elif md_field == ChunkMetadata.HEADINGS:
+                    result[md_field.name.lower()] = section.md.headings
+                elif md_field == ChunkMetadata.SOURCELINE:
+                    result[md_field.name.lower()] = section.md.source_line
+            return result
+
+    def _html_sections(self) -> Generator[DocumentSection, None, None]:
+        """Create DocumentSections reflecting the html-specific separators"""
+        import bs4
+        emit_on_paragraph = Separator.PARAGRAPH in self._separators or Separator.SENTENCE in self._separators
+        emit_on_heading = Separator.HEADING in self._separators or emit_on_paragraph
+        # current state
+        text_section = ''  # currently accumulated text
+        headings: Dict[int, str] = {}   # current state of observed headings (level -> text)
+        sourceline = 0  # most recently seen sourceline
+
+        def update_md(el: bs4.Tag) -> None:
+            # update current state
+            nonlocal headings, sourceline
+            sourceline = el.sourceline
+            if el.name in ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']:
+                level = int(el.name[1])
+                # remove the previously seen lower levels
+                lower_levels = [l for l in headings.keys() if l > level]
+                for l in lower_levels:
+                    del headings[l]
+                headings[level] = el.get_text().strip()
+
+        def emit() -> None:
+            nonlocal text_section, headings, sourceline
+            if len(text_section) > 0:
+                md = DocumentSectionMd(sourceline, headings.copy())
+                yield DocumentSection(text=text_section, md=md)
+                text_section = ''
+
+        def process_element(el: bs4.PageElement) -> Generator[DocumentSection, None, None]:
+            # process the element and emit sections as necessary
+            nonlocal text_section, headings, sourceline, emit_on_heading, emit_on_paragraph
+            if el.name in self._skip_tags:
+                return
+
+            if isinstance(el, bs4.NavigableString):
+                # accumulate text until we see a tag we care about
+                text = el.get_text().strip()
+                if len(text) > 0:
+                    text_section += ' ' + text
+                return
+
+            if el.name in ['h1', 'h2', 'h3', 'h4', 'h5', 'h6']:
+                if emit_on_heading:
+                    yield from emit()
+                update_md(el)
+            elif el.name == 'p':
+                if emit_on_paragraph:
+                    yield from emit()
+                update_md(el)
+            for child in el.children:
+                yield from process_element(child)
+
+        yield from process_element(self._doc_handle.bs_doc)
+        yield from emit()
+
+    def _markdown_sections(self) -> Generator[DocumentSection, None, None]:
+        """Create DocumentSections reflecting the html-specific separators"""
+        assert self._doc_handle.md_ast is not None
+        emit_on_paragraph = Separator.PARAGRAPH in self._separators or Separator.SENTENCE in self._separators
+        emit_on_heading = Separator.HEADING in self._separators or emit_on_paragraph
+        # current state
+        text_section = ''  # currently accumulated text
+        headings: Dict[int, str] = {}   # current state of observed headings (level -> text)
+
+        def update_headings(heading: Dict) -> None:
+            # update current state
+            nonlocal headings
+            assert 'type' in heading and heading['type'] == 'heading'
+            level = heading['attrs']['level']
+            text = heading['children'][0]['raw'].strip()
+            # remove the previously seen lower levels
+            lower_levels = [l for l in headings.keys() if l > level]
+            for l in lower_levels:
+                del headings[l]
+            headings[level] = text
+
+        def emit() -> None:
+            nonlocal text_section, headings
+            if len(text_section) > 0:
+                md = DocumentSectionMd(0, headings.copy())
+                yield DocumentSection(text=text_section, md=md)
+                text_section = ''
+
+        def process_element(el: Dict) -> Generator[DocumentSection, None, None]:
+            # process the element and emit sections as necessary
+            nonlocal text_section, headings, emit_on_heading, emit_on_paragraph
+            assert 'type' in el
+
+            if el['type'] == 'text':
+                # accumulate text until we see a separator element
+                text = el['raw'].strip()
+                if len(text) > 0:
+                    text_section += ' ' + text
+                return
+
+            if el['type'] == 'heading':
+                if emit_on_heading:
+                    yield from emit()
+                update_headings(el)
+            elif el['type'] == 'paragraph':
+                if emit_on_paragraph:
+                    yield from emit()
+            if 'children' not in el:
+                return
+            for child in el['children']:
+                yield from process_element(child)
+
+        for el in self._doc_handle.md_ast:
+            yield from process_element(el)
+        yield from emit()
+
+    def _sentence_sections(self, input_sections: Iterable[DocumentSection]) -> Generator[DocumentSection, None, None]:
+        """Split the input sections into sentences"""
+        for section in input_sections:
+            if section.text is not None:
+                doc = Env.get().spacy_nlp(section.text)
+                for sent in doc.sents:
+                    yield DocumentSection(text=sent.text, md=section.md)
+
+    def _token_chunks(self, input: Iterable[DocumentSection]) -> Generator[DocumentSection, None, None]:
+        import tiktoken
+        if self._tiktoken_target_model is not None:
+            encoding = tiktoken.encoding_for_model(self._tiktoken_target_model)
+        else:
+            encoding = tiktoken.get_encoding(self._tiktoken_encoding)
+        assert self._limit > 0 and self._overlap >= 0
+
+        for section in input:
+            if section.text is None:
+                continue
+            tokens = encoding.encode(section.text)
+            start_idx = 0
+            while start_idx < len(tokens):
+                end_idx = min(start_idx + self._limit, len(tokens))
+                text = encoding.decode(tokens[start_idx:end_idx])
+                yield DocumentSection(text=text, md=section.md)
+                start_idx += self._limit - self._overlap
+
+    def _char_chunks(self, input: Iterable[DocumentSection]) -> Generator[DocumentSection, None, None]:
+        for section in input:
+            if section.text is None:
+                continue
+            start_idx = 0
+            while start_idx < len(section.text):
+                end_idx = min(start_idx + self._limit, len(section.text))
+                text = section.text[start_idx:end_idx]
+                yield DocumentSection(text=text, md=section.md)
+                start_idx += self._limit - self._overlap
+
+    def close(self) -> None:
+        pass
+
+    def set_pos(self, pos: int) -> None:
+        pass

pixeltable/iterators/video.py

@@ -0,0 +1,89 @@
+from typing import Dict, Any, List, Tuple
+from pathlib import Path
+import math
+import logging
+
+import cv2
+import PIL.Image
+
+from .base import ComponentIterator
+
+from pixeltable.type_system import ColumnType, VideoType, ImageType, IntType, FloatType
+from pixeltable.exceptions import Error
+
+
+_logger = logging.getLogger('pixeltable')
+
+class FrameIterator(ComponentIterator):
+    def __init__(self, video: str, fps: float = 0.0):
+        video_path = Path(video)
+        assert video_path.exists() and video_path.is_file()
+        self.video_path = video_path
+        self.fps = fps
+        self.video_reader = cv2.VideoCapture(str(video_path))
+        if not self.video_reader.isOpened():
+            raise Error(f'Failed to open video: {video}')
+        video_fps = int(self.video_reader.get(cv2.CAP_PROP_FPS))
+        if fps > video_fps:
+            raise Error(f'Video {video}: requested fps ({fps}) exceeds that of the video ({video_fps})')
+        self.frame_freq = int(video_fps / fps) if fps > 0 else 1
+        num_video_frames = int(self.video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
+        if num_video_frames == 0:
+            raise Error(f'Video {video}: failed to get number of frames')
+        # ceil: round up to ensure we count frame 0
+        self.num_frames = math.ceil(num_video_frames / self.frame_freq) if fps > 0 else num_video_frames
+        _logger.debug(f'FrameIterator: path={self.video_path} fps={self.fps}')
+
+        self.next_frame_idx = 0
+
+    @classmethod
+    def input_schema(cls) -> Dict[str, ColumnType]:
+        return {
+            'video': VideoType(nullable=False),
+            'fps': FloatType()
+        }
+
+    @classmethod
+    def output_schema(cls, *args: Any, **kwargs: Any) -> Tuple[Dict[str, ColumnType], List[str]]:
+        return {
+            'frame_idx': IntType(),
+            'pos_msec': FloatType(),
+            'pos_frame': FloatType(),
+            'frame': ImageType(),
+        }, ['frame']
+
+    def __next__(self) -> Dict[str, Any]:
+        while True:
+            pos_msec = self.video_reader.get(cv2.CAP_PROP_POS_MSEC)
+            pos_frame = self.video_reader.get(cv2.CAP_PROP_POS_FRAMES)
+            status, img = self.video_reader.read()
+            if not status:
+                _logger.debug(f'releasing video reader for {self.video_path}')
+                self.video_reader.release()
+                self.video_reader = None
+                raise StopIteration
+            if pos_frame % self.frame_freq == 0:
+                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+                result = {
+                    'frame_idx': self.next_frame_idx,
+                    'pos_msec': pos_msec,
+                    'pos_frame': pos_frame,
+                    'frame': PIL.Image.fromarray(img),
+                }
+                self.next_frame_idx += 1
+                # frame_freq > 1: jumping to the target frame here with video_reader.set() is far slower than just
+                # skipping the unwanted frames
+                return result
+
+    def close(self) -> None:
+        if self.video_reader is not None:
+            self.video_reader.release()
+            self.video_reader = None
+
+    def set_pos(self, pos: int) -> None:
+        """Seek to frame idx"""
+        if pos == self.next_frame_idx:
+            return
+        _logger.debug(f'seeking to frame {pos}')
+        self.video_reader.set(cv2.CAP_PROP_POS_FRAMES, pos * self.frame_freq)
+        self.next_frame_idx = pos

pixeltable/metadata/__init__.py

@@ -0,0 +1,54 @@
+import dataclasses
+import importlib
+import os
+import pkgutil
+from typing import Callable, Dict
+
+import sqlalchemy as sql
+import sqlalchemy.orm as orm
+
+from .schema import SystemInfo, SystemInfoMd
+
+# current version of the metadata; this is incremented whenever the metadata schema changes
+VERSION = 12
+
+
+def create_system_info(engine: sql.engine.Engine) -> None:
+    """Create the systemmetadata record"""
+    system_md = SystemInfoMd(schema_version=VERSION)
+    record = SystemInfo(md=dataclasses.asdict(system_md))
+    with orm.Session(engine, future=True) as session:
+        session.add(record)
+        session.flush()
+        session.commit()
+
+# conversion functions for upgrading the metadata schema from one version to the following
+# key: old schema version
+converter_cbs: Dict[int, Callable[[sql.engine.Engine], None]] = {}
+
+def register_converter(version: int, cb: Callable[[sql.engine.Engine], None]) -> None:
+    global converter_cbs
+    converter_cbs[version] = cb
+
+# load all converter modules
+for _, modname, _ in pkgutil.iter_modules([os.path.dirname(__file__) + '/converters']):
+    importlib.import_module('pixeltable.metadata.converters.' + modname)
+
+def upgrade_md(engine: sql.engine.Engine) -> None:
+    """Upgrade the metadata schema to the current version"""
+    with orm.Session(engine, future=True) as session:
+        system_info = session.query(SystemInfo).one().md
+        md_version = system_info['schema_version']
+        if md_version == VERSION:
+                return
+        while md_version < VERSION:
+            if md_version not in converter_cbs:
+                raise RuntimeError(f'No metadata converter for version {md_version}')
+            print(f'Converting metadata from version {md_version} to {md_version + 1}')
+            converter_cbs[md_version](engine)
+            md_version += 1
+        # update system info
+        conn = session.connection()
+        system_info_md = SystemInfoMd(schema_version=VERSION)
+        conn.execute(SystemInfo.__table__.update().values(md=dataclasses.asdict(system_info_md)))
+        session.commit()

pixeltable/metadata/converters/convert_10.py

@@ -0,0 +1,18 @@
+import sqlalchemy as sql
+
+from pixeltable.metadata.schema import Table, TableSchemaVersion
+from pixeltable.metadata import register_converter
+
+
+def convert_10(engine: sql.engine.Engine) -> None:
+    default_table_attrs = {"comment": None, "num_retained_versions": 10}
+    with engine.begin() as conn:
+        # Because `parameters` wasn't actually used for anything,
+        # we can simply delete it without any data loss.
+        conn.execute(sql.update(Table).values(md=Table.md - 'parameters'))
+        # Add `table_attrs` to all instances of tableschemaversions.md.
+        conn.execute(sql.update(TableSchemaVersion).values(md=TableSchemaVersion.md.concat(default_table_attrs)))
+    return
+
+
+register_converter(10, convert_10)