datachain 0.6.11__tar.gz → 0.7.1__tar.gz

{datachain-0.6.11 → datachain-0.7.1}/.github/workflows/tests.yml

@@ -104,7 +104,7 @@ jobs:
         shell: bash
 
       - name: Upload coverage report
-        uses: codecov/codecov-action@v4
+        uses: codecov/codecov-action@v5
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
           files: coverage.xml

{datachain-0.6.11 → datachain-0.7.1}/.pre-commit-config.yaml

@@ -24,7 +24,7 @@ repos:
       - id: trailing-whitespace
         exclude: '^LICENSES/'
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.7.3'
+    rev: 'v0.7.4'
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{datachain-0.6.11/src/datachain.egg-info → datachain-0.7.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.6.11
+Version: 0.7.1
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0
@@ -98,6 +98,7 @@ Requires-Dist: unstructured[embed-huggingface,pdf]<0.16.0; extra == "examples"
 Requires-Dist: pdfplumber==0.11.4; extra == "examples"
 Requires-Dist: huggingface_hub[hf_transfer]; extra == "examples"
 Requires-Dist: onnx==1.16.1; extra == "examples"
+Requires-Dist: ultralytics==8.3.29; extra == "examples"
 
 ================
 |logo| DataChain

{datachain-0.6.11 → datachain-0.7.1}/examples/computer_vision/openimage-detect.py

@@ -1,19 +1,11 @@
 import json
 
 from PIL import Image
-from pydantic import BaseModel
 
-from datachain import C, DataChain, File
+from datachain import C, DataChain, File, model
 from datachain.sql.functions import path
 
 
-class BBox(BaseModel):
-    x_min: int
-    x_max: int
-    y_min: int
-    y_max: int
-
-
 def openimage_detect(args):
     if len(args) != 2:
         raise ValueError("Group jpg-json mismatch")
@@ -30,11 +22,13 @@ def openimage_detect(args):
         detections = json.load(stream_json).get("detections", [])
 
     for i, detect in enumerate(detections):
-        bbox = BBox(
-            x_min=int(detect["XMin"] * img.width),
-            x_max=int(detect["XMax"] * img.width),
-            y_min=int(detect["YMin"] * img.height),
-            y_max=int(detect["YMax"] * img.height),
+        bbox = model.BBox.from_list(
+            [
+                detect["XMin"] * img.width,
+                detect["XMax"] * img.width,
+                detect["YMin"] * img.height,
+                detect["YMax"] * img.height,
+            ]
         )
 
         fstream = File(
@@ -56,7 +50,7 @@ source = "gs://datachain-demo/openimages-v6-test-jsonpairs/"
         openimage_detect,
         partition_by=path.file_stem(C("file.path")),
         params=["file"],
-        output={"file": File, "bbox": BBox},
+        output={"file": File, "bbox": model.BBox},
     )
     .show()
 )

datachain-0.7.1/examples/computer_vision/ultralytics-bbox.py

@@ -0,0 +1,22 @@
+from io import BytesIO
+
+from PIL import Image
+from ultralytics import YOLO
+
+from datachain import C, DataChain, File
+from datachain.model.ultralytics import YoloBBoxes
+
+
+def process_bboxes(yolo: YOLO, file: File) -> YoloBBoxes:
+    results = yolo(Image.open(BytesIO(file.read())))
+    return YoloBBoxes.from_results(results)
+
+
+(
+    DataChain.from_storage("gs://datachain-demo/openimages-v6-test-jsonpairs/")
+    .filter(C("file.path").glob("*.jpg"))
+    .limit(20)
+    .setup(yolo=lambda: YOLO("yolo11n.pt"))
+    .map(boxes=process_bboxes)
+    .show()
+)

datachain-0.7.1/examples/computer_vision/ultralytics-pose.py

@@ -0,0 +1,22 @@
+from io import BytesIO
+
+from PIL import Image
+from ultralytics import YOLO
+
+from datachain import C, DataChain, File
+from datachain.model.ultralytics import YoloPoses
+
+
+def process_poses(yolo: YOLO, file: File) -> YoloPoses:
+    results = yolo(Image.open(BytesIO(file.read())))
+    return YoloPoses.from_results(results)
+
+
+(
+    DataChain.from_storage("gs://datachain-demo/openimages-v6-test-jsonpairs/")
+    .filter(C("file.path").glob("*.jpg"))
+    .limit(20)
+    .setup(yolo=lambda: YOLO("yolo11n-pose.pt"))
+    .map(poses=process_poses)
+    .show()
+)

datachain-0.7.1/examples/computer_vision/ultralytics-segment.py

@@ -0,0 +1,22 @@
+from io import BytesIO
+
+from PIL import Image
+from ultralytics import YOLO
+
+from datachain import C, DataChain, File
+from datachain.model.ultralytics import YoloSegments
+
+
+def process_segments(yolo: YOLO, file: File) -> YoloSegments:
+    results = yolo(Image.open(BytesIO(file.read())))
+    return YoloSegments.from_results(results)
+
+
+(
+    DataChain.from_storage("gs://datachain-demo/openimages-v6-test-jsonpairs/")
+    .filter(C("file.path").glob("*.jpg"))
+    .limit(20)
+    .setup(yolo=lambda: YOLO("yolo11n-seg.pt"))
+    .map(segments=process_segments)
+    .show()
+)

{datachain-0.6.11 → datachain-0.7.1}/pyproject.toml

@@ -110,7 +110,8 @@ examples = [
   "unstructured[pdf,embed-huggingface]<0.16.0",
   "pdfplumber==0.11.4",
   "huggingface_hub[hf_transfer]",
-  "onnx==1.16.1"
+  "onnx==1.16.1",
+  "ultralytics==8.3.29"
 ]
 
 [project.urls]

{datachain-0.6.11 → datachain-0.7.1}/src/datachain/__init__.py

@@ -1,4 +1,4 @@
-from datachain.lib import func, models
+from datachain.lib import func
 from datachain.lib.data_model import DataModel, DataType, is_chain_type
 from datachain.lib.dc import C, Column, DataChain, Sys
 from datachain.lib.file import (
@@ -38,6 +38,5 @@ __all__ = [
     "func",
     "is_chain_type",
     "metrics",
-    "models",
     "param",
 ]

{datachain-0.6.11 → datachain-0.7.1}/src/datachain/asyn.py

@@ -1,5 +1,13 @@
 import asyncio
-from collections.abc import AsyncIterable, Awaitable, Coroutine, Iterable, Iterator
+import threading
+from collections.abc import (
+    AsyncIterable,
+    Awaitable,
+    Coroutine,
+    Generator,
+    Iterable,
+    Iterator,
+)
 from concurrent.futures import ThreadPoolExecutor
 from heapq import heappop, heappush
 from typing import Any, Callable, Generic, Optional, TypeVar
@@ -47,6 +55,7 @@ class AsyncMapper(Generic[InputT, ResultT]):
         self.loop = get_loop() if loop is None else loop
         self.pool = ThreadPoolExecutor(workers)
         self._tasks: set[asyncio.Task] = set()
+        self._shutdown_producer = threading.Event()
 
     def start_task(self, coro: Coroutine) -> asyncio.Task:
         task = self.loop.create_task(coro)
@@ -54,9 +63,31 @@ class AsyncMapper(Generic[InputT, ResultT]):
         task.add_done_callback(self._tasks.discard)
         return task
 
-    async def produce(self) -> None:
+    def _produce(self) -> None:
         for item in self.iterable:
-            await self.work_queue.put(item)
+            if self._shutdown_producer.is_set():
+                return
+            fut = asyncio.run_coroutine_threadsafe(self.work_queue.put(item), self.loop)
+            fut.result()  # wait until the item is in the queue
+
+    async def produce(self) -> None:
+        await self.to_thread(self._produce)
+
+    def shutdown_producer(self) -> None:
+        """
+        Signal the producer to stop and drain any remaining items from the work_queue.
+
+        This method sets an internal event, `_shutdown_producer`, which tells the
+        producer that it should stop adding items to the queue. To ensure that the
+        producer notices this signal promptly, we also attempt to drain any items
+        currently in the queue, clearing it so that the event can be checked without
+        delay.
+        """
+        self._shutdown_producer.set()
+        q = self.work_queue
+        while not q.empty():
+            q.get_nowait()
+            q.task_done()
 
     async def worker(self) -> None:
         while (item := await self.work_queue.get()) is not None:
@@ -132,7 +163,7 @@ class AsyncMapper(Generic[InputT, ResultT]):
             self.result_queue.get_nowait()
         await self.result_queue.put(None)
 
-    def iterate(self, timeout=None) -> Iterable[ResultT]:
+    def iterate(self, timeout=None) -> Generator[ResultT, None, None]:
         init = asyncio.run_coroutine_threadsafe(self.init(), self.loop)
         init.result(timeout=1)
         async_run = asyncio.run_coroutine_threadsafe(self.run(), self.loop)
@@ -145,6 +176,7 @@ class AsyncMapper(Generic[InputT, ResultT]):
             if exc := async_run.exception():
                 raise exc
         finally:
+            self.shutdown_producer()
             if not async_run.done():
                 async_run.cancel()
 

{datachain-0.6.11 → datachain-0.7.1}/src/datachain/data_storage/warehouse.py

@@ -232,7 +232,10 @@ class AbstractWarehouse(ABC, Serializable):
                 if limit < page_size:
                     paginated_query = paginated_query.limit(None).limit(limit)
 
-            results = self.dataset_rows_select(paginated_query.offset(offset))
+            # Ensure we're using a thread-local connection
+            with self.clone() as wh:
+                # Cursor results are not thread-safe, so we convert them to a list
+                results = list(wh.dataset_rows_select(paginated_query.offset(offset)))
 
             processed = False
             for row in results:

{datachain-0.6.11 → datachain-0.7.1}/src/datachain/lib/dc.py

@@ -334,6 +334,7 @@ class DataChain:
         parallel=None,
         workers=None,
         min_task_size=None,
+        prefetch: Optional[int] = None,
         sys: Optional[bool] = None,
     ) -> "Self":
         """Change settings for chain.
@@ -360,7 +361,7 @@ class DataChain:
         if sys is None:
             sys = self._sys
         settings = copy.copy(self._settings)
-        settings.add(Settings(cache, parallel, workers, min_task_size))
+        settings.add(Settings(cache, parallel, workers, min_task_size, prefetch))
         return self._evolve(settings=settings, _sys=sys)
 
     def reset_settings(self, settings: Optional[Settings] = None) -> "Self":
@@ -882,6 +883,8 @@ class DataChain:
             ```
         """
         udf_obj = self._udf_to_obj(Mapper, func, params, output, signal_map)
+        if (prefetch := self._settings.prefetch) is not None:
+            udf_obj.prefetch = prefetch
 
         return self._evolve(
             query=self._query.add_signals(
@@ -919,6 +922,8 @@ class DataChain:
             ```
         """
         udf_obj = self._udf_to_obj(Generator, func, params, output, signal_map)
+        if (prefetch := self._settings.prefetch) is not None:
+            udf_obj.prefetch = prefetch
         return self._evolve(
             query=self._query.generate(
                 udf_obj.to_udf_wrapper(),

{datachain-0.6.11 → datachain-0.7.1}/src/datachain/lib/file.py

@@ -268,6 +268,11 @@ class File(DataModel):
         client = self._catalog.get_client(self.source)
         client.download(self, callback=self._download_cb)
 
+    async def _prefetch(self) -> None:
+        if self._caching_enabled:
+            client = self._catalog.get_client(self.source)
+            await client._download(self, callback=self._download_cb)
+
     def get_local_path(self) -> Optional[str]:
         """Return path to a file in a local cache.
 

{datachain-0.6.11 → datachain-0.7.1}/src/datachain/lib/settings.py

@@ -7,11 +7,19 @@ class SettingsError(DataChainParamsError):
 
 
 class Settings:
-    def __init__(self, cache=None, parallel=None, workers=None, min_task_size=None):
+    def __init__(
+        self,
+        cache=None,
+        parallel=None,
+        workers=None,
+        min_task_size=None,
+        prefetch=None,
+    ):
         self._cache = cache
         self.parallel = parallel
         self._workers = workers
         self.min_task_size = min_task_size
+        self.prefetch = prefetch
 
         if not isinstance(cache, bool) and cache is not None:
             raise SettingsError(
@@ -66,3 +74,5 @@ class Settings:
         self.parallel = settings.parallel or self.parallel
         self._workers = settings._workers or self._workers
         self.min_task_size = settings.min_task_size or self.min_task_size
+        if settings.prefetch is not None:
+            self.prefetch = settings.prefetch

{datachain-0.6.11 → datachain-0.7.1}/src/datachain/lib/udf.py

@@ -1,3 +1,4 @@
+import contextlib
 import sys
 import traceback
 from collections.abc import Iterable, Iterator, Mapping, Sequence
@@ -7,6 +8,7 @@ import attrs
 from fsspec.callbacks import DEFAULT_CALLBACK, Callback
 from pydantic import BaseModel
 
+from datachain.asyn import AsyncMapper
 from datachain.dataset import RowDict
 from datachain.lib.convert.flatten import flatten
 from datachain.lib.data_model import DataValue
@@ -21,6 +23,8 @@ from datachain.query.batch import (
 )
 
 if TYPE_CHECKING:
+    from collections import abc
+
     from typing_extensions import Self
 
     from datachain.catalog import Catalog
@@ -276,9 +280,18 @@ class UDFBase(AbstractUDF):
         return result_objs
 
 
+async def _prefetch_input(row):
+    for obj in row:
+        if isinstance(obj, File):
+            await obj._prefetch()
+    return row
+
+
 class Mapper(UDFBase):
     """Inherit from this class to pass to `DataChain.map()`."""
 
+    prefetch: int = 2
+
     def run(
         self,
         udf_fields: "Sequence[str]",
@@ -290,16 +303,22 @@ class Mapper(UDFBase):
     ) -> Iterator[Iterable[UDFResult]]:
         self.catalog = catalog
         self.setup()
-
-        for row in udf_inputs:
-            id_, *udf_args = self._prepare_row_and_id(
-                row, udf_fields, cache, download_cb
-            )
-            result_objs = self.process_safe(udf_args)
-            udf_output = self._flatten_row(result_objs)
-            output = [{"sys__id": id_} | dict(zip(self.signal_names, udf_output))]
-            processed_cb.relative_update(1)
-            yield output
+        prepared_inputs: abc.Generator[Sequence[Any], None, None] = (
+            self._prepare_row_and_id(row, udf_fields, cache, download_cb)
+            for row in udf_inputs
+        )
+        if self.prefetch > 0:
+            prepared_inputs = AsyncMapper(
+                _prefetch_input, prepared_inputs, workers=self.prefetch
+            ).iterate()
+
+        with contextlib.closing(prepared_inputs):
+            for id_, *udf_args in prepared_inputs:
+                result_objs = self.process_safe(udf_args)
+                udf_output = self._flatten_row(result_objs)
+                output = [{"sys__id": id_} | dict(zip(self.signal_names, udf_output))]
+                processed_cb.relative_update(1)
+                yield output
 
         self.teardown()
 
@@ -349,6 +368,7 @@ class Generator(UDFBase):
     """Inherit from this class to pass to `DataChain.gen()`."""
 
     is_output_batched = True
+    prefetch: int = 2
 
     def run(
         self,
@@ -361,14 +381,21 @@ class Generator(UDFBase):
     ) -> Iterator[Iterable[UDFResult]]:
         self.catalog = catalog
         self.setup()
-
-        for row in udf_inputs:
-            udf_args = self._prepare_row(row, udf_fields, cache, download_cb)
-            result_objs = self.process_safe(udf_args)
-            udf_outputs = (self._flatten_row(row) for row in result_objs)
-            output = (dict(zip(self.signal_names, row)) for row in udf_outputs)
-            processed_cb.relative_update(1)
-            yield output
+        prepared_inputs: abc.Generator[Sequence[Any], None, None] = (
+            self._prepare_row(row, udf_fields, cache, download_cb) for row in udf_inputs
+        )
+        if self.prefetch > 0:
+            prepared_inputs = AsyncMapper(
+                _prefetch_input, prepared_inputs, workers=self.prefetch
+            ).iterate()
+
+        with contextlib.closing(prepared_inputs):
+            for row in prepared_inputs:
+                result_objs = self.process_safe(row)
+                udf_outputs = (self._flatten_row(row) for row in result_objs)
+                output = (dict(zip(self.signal_names, row)) for row in udf_outputs)
+                processed_cb.relative_update(1)
+                yield output
 
         self.teardown()
 

datachain-0.7.1/src/datachain/model/__init__.py

@@ -0,0 +1,6 @@
+from . import ultralytics
+from .bbox import BBox, OBBox
+from .pose import Pose, Pose3D
+from .segment import Segment
+
+__all__ = ["BBox", "OBBox", "Pose", "Pose3D", "Segment", "ultralytics"]

datachain-0.7.1/src/datachain/model/bbox.py

@@ -0,0 +1,102 @@
+from pydantic import Field
+
+from datachain.lib.data_model import DataModel
+
+
+class BBox(DataModel):
+    """
+    A data model for representing bounding box.
+
+    Attributes:
+        title (str): The title of the bounding box.
+        coords (list[int]): The coordinates of the bounding box.
+
+    The bounding box is defined by two points:
+        - (x1, y1): The top-left corner of the box.
+        - (x2, y2): The bottom-right corner of the box.
+    """
+
+    title: str = Field(default="")
+    coords: list[int] = Field(default=None)
+
+    @staticmethod
+    def from_list(coords: list[float], title: str = "") -> "BBox":
+        assert len(coords) == 4, "Bounding box must be a list of 4 coordinates."
+        assert all(
+            isinstance(value, (int, float)) for value in coords
+        ), "Bounding box coordinates must be floats or integers."
+        return BBox(
+            title=title,
+            coords=[round(c) for c in coords],
+        )
+
+    @staticmethod
+    def from_dict(coords: dict[str, float], title: str = "") -> "BBox":
+        assert isinstance(coords, dict) and set(coords) == {
+            "x1",
+            "y1",
+            "x2",
+            "y2",
+        }, "Bounding box must be a dictionary with keys 'x1', 'y1', 'x2' and 'y2'."
+        return BBox.from_list(
+            [coords["x1"], coords["y1"], coords["x2"], coords["y2"]],
+            title=title,
+        )
+
+
+class OBBox(DataModel):
+    """
+    A data model for representing oriented bounding boxes.
+
+    Attributes:
+        title (str): The title of the oriented bounding box.
+        coords (list[int]): The coordinates of the oriented bounding box.
+
+    The oriented bounding box is defined by four points:
+        - (x1, y1): The first corner of the box.
+        - (x2, y2): The second corner of the box.
+        - (x3, y3): The third corner of the box.
+        - (x4, y4): The fourth corner of the box.
+    """
+
+    title: str = Field(default="")
+    coords: list[int] = Field(default=None)
+
+    @staticmethod
+    def from_list(coords: list[float], title: str = "") -> "OBBox":
+        assert (
+            len(coords) == 8
+        ), "Oriented bounding box must be a list of 8 coordinates."
+        assert all(
+            isinstance(value, (int, float)) for value in coords
+        ), "Oriented bounding box coordinates must be floats or integers."
+        return OBBox(
+            title=title,
+            coords=[round(c) for c in coords],
+        )
+
+    @staticmethod
+    def from_dict(coords: dict[str, float], title: str = "") -> "OBBox":
+        assert isinstance(coords, dict) and set(coords) == {
+            "x1",
+            "y1",
+            "x2",
+            "y2",
+            "x3",
+            "y3",
+            "x4",
+            "y4",
+        }, "Oriented bounding box must be a dictionary with coordinates."
+        return OBBox.from_list(
+            [
+                coords["x1"],
+                coords["y1"],
+                coords["x2"],
+                coords["y2"],
+                coords["x3"],
+                coords["y3"],
+                coords["x4"],
+                coords["y4"],
+            ],
+            title=title,
+        )

datachain-0.7.1/src/datachain/model/pose.py

@@ -0,0 +1,88 @@
+from pydantic import Field
+
+from datachain.lib.data_model import DataModel
+
+
+class Pose(DataModel):
+    """
+    A data model for representing pose keypoints.
+
+    Attributes:
+        x (list[int]): The x-coordinates of the keypoints.
+        y (list[int]): The y-coordinates of the keypoints.
+
+    The keypoints are represented as lists of x and y coordinates, where each index
+    corresponds to a specific body part.
+    """
+
+    x: list[int] = Field(default=None)
+    y: list[int] = Field(default=None)
+
+    @staticmethod
+    def from_list(points: list[list[float]]) -> "Pose":
+        assert len(points) == 2, "Pose must be a list of 2 lists: x and y coordinates."
+        points_x, points_y = points
+        assert (
+            len(points_x) == len(points_y) == 17
+        ), "Pose x and y coordinates must have the same length of 17."
+        assert all(
+            isinstance(value, (int, float)) for value in [*points_x, *points_y]
+        ), "Pose coordinates must be floats or integers."
+        return Pose(
+            x=[round(coord) for coord in points_x],
+            y=[round(coord) for coord in points_y],
+        )
+
+    @staticmethod
+    def from_dict(points: dict[str, list[float]]) -> "Pose":
+        assert isinstance(points, dict) and set(points) == {
+            "x",
+            "y",
+        }, "Pose must be a dict with keys 'x' and 'y'."
+        return Pose.from_list([points["x"], points["y"]])
+
+
+class Pose3D(DataModel):
+    """
+    A data model for representing 3D pose keypoints.
+
+    Attributes:
+        x (list[int]): The x-coordinates of the keypoints.
+        y (list[int]): The y-coordinates of the keypoints.
+        visible (list[float]): The visibility of the keypoints.
+
+    The keypoints are represented as lists of x, y, and visibility values,
+    where each index corresponds to a specific body part.
+    """
+
+    x: list[int] = Field(default=None)
+    y: list[int] = Field(default=None)
+    visible: list[float] = Field(default=None)
+
+    @staticmethod
+    def from_list(points: list[list[float]]) -> "Pose3D":
+        assert (
+            len(points) == 3
+        ), "Pose3D must be a list of 3 lists: x, y coordinates and visible."
+        points_x, points_y, points_v = points
+        assert (
+            len(points_x) == len(points_y) == len(points_v) == 17
+        ), "Pose3D x, y coordinates and visible must have the same length of 17."
+        assert all(
+            isinstance(value, (int, float))
+            for value in [*points_x, *points_y, *points_v]
+        ), "Pose3D coordinates must be floats or integers."
+        return Pose3D(
+            x=[round(coord) for coord in points_x],
+            y=[round(coord) for coord in points_y],
+            visible=points_v,
+        )
+
+    @staticmethod
+    def from_dict(points: dict[str, list[float]]) -> "Pose3D":
+        assert isinstance(points, dict) and set(points) == {
+            "x",
+            "y",
+            "visible",
+        }, "Pose3D must be a dict with keys 'x', 'y' and 'visible'."
+        return Pose3D.from_list([points["x"], points["y"], points["visible"]])