label-studio-sdk 0.0.34__py3-none-any.whl → 1.0.1__py3-none-any.whl

label_studio_sdk/_extensions/label_studio_tools/etl/beam.py

@@ -0,0 +1,34 @@
+import apache_beam as beam
+from registry import dataloader, call_dataloader
+
+
+class CallDataLoader(beam.DoFn):
+    def __init__(self, dataloader_name):
+        self.dataloader_name = dataloader_name
+
+    def setup(self):
+        print("Setup!")
+        pass
+
+    def process(self, element, *args, **kwargs):
+        return call_dataloader(self.dataloader_name, element, *args, **kwargs)
+
+    def teardown(self):
+        print("Teardown!")
+        pass
+
+
+@dataloader("test")
+def f(i, k):
+    print("I!!", k)
+    yield i
+
+
+with beam.Pipeline() as pipeline:
+    results = (
+        pipeline
+        | "Get storage keys" >> beam.io.gcp.gcsio.list_files("bucket")
+        | "Run dataloader"
+        >> beam.FlatMap(fn=CallDataLoader("test"), **dataloader_kwargs)
+        | "Load into Vector DB" >> beam.ParDo(fn=LoadDataRecordInDB)
+    )

label_studio_sdk/_extensions/label_studio_tools/etl/example.py

@@ -0,0 +1,17 @@
+from label_studio_sdk._extensions.label_studio_tools.etl.registry import (
+    transform,
+    call_dataloader,
+)
+
+
+@transform.datarecords("read_base64_from_blob")
+def gcs_blob_to_base64(
+    key: str,
+    bucket_name: str,
+    google_project_id: str = None,
+):
+    print(key, bucket_name)
+
+
+if __name__ == "__main__":
+    call_dataloader("gcs_blob_reader", key=12, bucket_name=34)

label_studio_sdk/_extensions/label_studio_tools/etl/registry.py

@@ -0,0 +1,67 @@
+import logging
+
+from functools import wraps
+from typing import Callable, Iterable, Dict, Any
+
+logger = logging.getLogger(__name__)
+
+
+class extract:
+    pass
+
+
+class transform:
+    DataRecords = Iterable[Dict[str, Any]]
+    DataRecordsProducer = Callable[..., DataRecords]
+
+    class NameAlreadyExists(KeyError):
+        pass
+
+    class NameNotFound(NotImplementedError):
+        pass
+
+    class WrongOutput(ValueError):
+        pass
+
+    registry = {}
+
+    @staticmethod
+    def datarecords(name: str = None):
+        """Class method to register transform function to produce DataRecords
+        with specified name to the internal registry.
+        Can output any iterable for dicts (generators, iterators, list, tuples)
+        For example:
+
+        @transform.datarecords("DataRecords from GCS file")
+        def read_image(bucket_name, key):
+            yield {"image": blob.download_as_string()}
+        """
+
+        def register_transform_func(
+            transform_func: transform.DataRecordsProducer,
+        ) -> transform.DataRecordsProducer:
+            if name in transform.registry:
+                raise transform.NameAlreadyExists(
+                    f"Dataloader {name} already exists. Will replace it"
+                )
+            transform.registry[name] = transform_func
+
+            @wraps(transform_func)
+            def inner_call_transform_func(*args, **kwargs) -> transform.DataRecords:
+                data_records = transform_func(*args, **kwargs)
+                # TODO: assert Data record type in a proper way, typing.assert_type is supported in Python 3.11
+                return data_records
+
+            return inner_call_transform_func
+
+        return register_transform_func
+
+    @staticmethod
+    def invoke(name, *args, **kwargs):
+        if name not in transform.registry:
+            raise transform.NameNotFound(name)
+        return transform.registry[name](*args, **kwargs)
+
+
+class load:
+    pass

label_studio_sdk/_extensions/label_studio_tools/postprocessing/video.py

@@ -0,0 +1,97 @@
+from copy import deepcopy
+from operator import itemgetter
+
+from label_studio_sdk._extensions.label_studio_tools.core.label_config import (
+    _VIDEO_TRACKING_TAGS,
+)
+
+
+def extract_key_frames(results):
+    """
+    Extract frames from key frames from video annotation results
+    :param results: Annotation results
+    :return: Frames in LS format
+    """
+    final_results = []
+    for result in results:
+        temp = deepcopy(result)
+        if result["type"].lower() not in _VIDEO_TRACKING_TAGS:
+            final_results.append(result)
+            continue
+        sequence = result["value"]["sequence"]
+        if len(sequence) < 1:
+            continue
+        sequence = sorted(sequence, key=lambda d: d["frame"])
+        exclude_first = False
+        for i in range(len(sequence)):
+            frame_a = sequence[i]
+            frame_b = {} if i == len(sequence) - 1 else sequence[i + 1]
+            temp["value"]["sequence"].extend(
+                _construct_result_from_frames(
+                    frame1=frame_a,
+                    frame2=frame_b,
+                    frameCount=result["value"].get("framesCount", 0),
+                    duration=result["value"].get("duration", 0),
+                    exclude_first=exclude_first,
+                )
+            )
+            exclude_first = frame_a["enabled"]
+        temp["value"]["sequence"] = sorted(
+            temp["value"]["sequence"], key=itemgetter("frame")
+        )
+        final_results.append(temp)
+    return final_results
+
+
+def _construct_result_from_frames(
+    frame1, frame2, frameCount=0, duration=0, exclude_first=True
+):
+    """
+    Construct frames between 2 keyframes
+    :param frame1: First frame in sequence
+    :param frame2: Next frame in sequence
+    :param frameCount: Total frame count in the video
+    :param duration: Total duration of the video
+    :param exclude_first: Exclude first result to deduplicate results
+    :return: List of frames
+    """
+    final_results = []
+    if not frame1["enabled"]:
+        return []
+    if len(frame2) > 0:
+        if frame1["frame"] > frame2["frame"]:
+            return []
+        frame_count = frame2["frame"] - frame1["frame"] + 1
+    else:
+        frame_count = frameCount - frame1["frame"] + 1
+    start_i = 1 if exclude_first else 0
+    for i in range(start_i, frame_count):
+        frame_number = i + frame1["frame"]
+        delta = i / max(1, (frame_count - 1))
+        deltas = {}
+        for v in ["x", "y", "rotation", "width", "height", "time"]:
+            deltas[v] = (
+                0
+                if (frame1[v] == frame2.get(v) or not frame2)
+                else (frame2.get(v, 0) - frame1[v]) * delta
+            )
+        result = deepcopy(frame1)
+        result.update(
+            {
+                "x": frame1["x"] + deltas["x"],
+                "y": frame1["y"] + deltas["y"],
+                "width": frame1["width"] + deltas["width"],
+                "height": frame1["height"] + deltas["height"],
+                "rotation": frame1["rotation"] + deltas["rotation"],
+                "frame": frame_number,
+                "time": round(frame1["time"] + deltas["time"], 2),
+            }
+        )
+        if frame_number not in [frame1.get("frame"), frame2.get("frame")]:
+            result["auto"] = True
+            if deltas["time"] == 0 and duration > 0:
+                result["time"] = round(
+                    frame1["time"] + delta * (duration - frame1["time"]), 2
+                )
+            final_results.append(result)
+    return final_results

label_studio_sdk/_extensions/pager_ext.py

@@ -0,0 +1,49 @@
+import typing
+from label_studio_sdk.core.pagination import SyncPager, AsyncPager, T
+from label_studio_sdk.core.api_error import ApiError
+
+# This is a custom extension of the autogenerated SyncPager and AsyncPager classes
+# that works with the Label Studio SDK's default pagination behavior
+# that throws 404 errors at the end of the pagination.
+
+
+class SyncPagerExt(SyncPager, typing.Generic[T]):
+
+    @classmethod
+    def from_sync_pager(cls, sync_pager: SyncPager) -> 'SyncPagerExt':
+        return cls(
+            get_next=sync_pager.get_next,
+            has_next=sync_pager.has_next,
+            items=sync_pager.items
+        )
+
+    def __iter__(self) -> typing.Iterator[T]:  # type: ignore
+        # Extends the iterator to catch 404 errors at the end of the pagination
+        try:
+            for item in super().__iter__():
+                yield item
+        except ApiError as exc:
+            if exc.status_code == 404:
+                return
+            raise
+
+
+class AsyncPagerExt(AsyncPager, typing.Generic[T]):
+
+    @classmethod
+    async def from_async_pager(cls, async_pager: AsyncPager) -> 'AsyncPagerExt':
+        return cls(
+            get_next=async_pager.get_next,
+            has_next=async_pager.has_next,
+            items=async_pager.items
+        )
+
+    async def __aiter__(self) -> typing.AsyncIterator[T]:  # type: ignore
+        # Extends the iterator to catch 404 errors at the end of the pagination
+        try:
+            async for item in super().__aiter__():
+                yield item
+        except ApiError as exc:
+            if exc.status_code == 404:
+                return
+            raise

label_studio_sdk/_legacy/__init__.py

@@ -0,0 +1,11 @@
+""" .. include::../docs/index.md
+"""
+
+from .client import Client
+from .project import Project
+from .utils import parse_config
+
+__pdoc__ = {"label_interface": False}
+
+
+__version__ = '0.0.35.dev'