datachain 0.6.10__py3-none-any.whl → 0.7.0__py3-none-any.whl

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/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/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/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/lib/models/__init__.py

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

datachain/lib/models/bbox.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 from pydantic import Field
 
 from datachain.lib.data_model import DataModel
@@ -9,7 +11,10 @@ class BBox(DataModel):
 
     Attributes:
         title (str): The title of the bounding box.
-        coords (list[int]): The coordinates of the bounding box.
+        x1 (float): The x-coordinate of the top-left corner of the bounding box.
+        y1 (float): The y-coordinate of the top-left corner of the bounding box.
+        x2 (float): The x-coordinate of the bottom-right corner of the bounding box.
+        y2 (float): The y-coordinate of the bottom-right corner of the bounding box.
 
     The bounding box is defined by two points:
         - (x1, y1): The top-left corner of the box.
@@ -17,100 +22,24 @@ class BBox(DataModel):
     """
 
     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 coordinates must be a list of 4 floats."
-        assert all(
-            isinstance(value, (int, float)) for value in coords
-        ), "Bounding box coordinates must be integers or floats."
-        return BBox(
-            title=title,
-            coords=[round(c) for c in coords],
-        )
-
-    @staticmethod
-    def from_dict(coords: dict[str, float], title: str = "") -> "BBox":
-        assert (
-            len(coords) == 4
-        ), "Bounding box coordinates must be a dictionary of 4 floats."
-        assert set(coords) == {
-            "x1",
-            "y1",
-            "x2",
-            "y2",
-        }, "Bounding box coordinates must contain keys with coordinates."
-        assert all(
-            isinstance(value, (int, float)) for value in coords.values()
-        ), "Bounding box coordinates must be integers or floats."
-        return BBox(
-            title=title,
-            coords=[
-                round(coords["x1"]),
-                round(coords["y1"]),
-                round(coords["x2"]),
-                round(coords["y2"]),
-            ],
-        )
-
-
-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 coordinates must be a list of 8 floats."
-        assert all(
-            isinstance(value, (int, float)) for value in coords
-        ), "Oriented bounding box coordinates must be integers or floats."
-        return OBBox(
-            title=title,
-            coords=[round(c) for c in coords],
-        )
+    x1: float = Field(default=0)
+    y1: float = Field(default=0)
+    x2: float = Field(default=0)
+    y2: float = Field(default=0)
 
     @staticmethod
-    def from_dict(coords: dict[str, float], title: str = "") -> "OBBox":
-        assert set(coords) == {
-            "x1",
-            "y1",
-            "x2",
-            "y2",
-            "x3",
-            "y3",
-            "x4",
-            "y4",
-        }, "Oriented bounding box coordinates must contain keys with coordinates."
-        assert all(
-            isinstance(value, (int, float)) for value in coords.values()
-        ), "Oriented bounding box coordinates must be integers or floats."
-        return OBBox(
-            title=title,
-            coords=[
-                round(coords["x1"]),
-                round(coords["y1"]),
-                round(coords["x2"]),
-                round(coords["y2"]),
-                round(coords["x3"]),
-                round(coords["y3"]),
-                round(coords["x4"]),
-                round(coords["y4"]),
-            ],
-        )
+    def from_xywh(bbox: list[float], title: Optional[str] = None) -> "BBox":
+        """
+        Converts a bounding box in (x, y, width, height) format
+        to a BBox data model instance.
+
+        Args:
+            bbox (list[float]): A bounding box, represented as a list
+                                of four floats [x, y, width, height].
+
+        Returns:
+            BBox2D: An instance of the BBox data model.
+        """
+        assert len(bbox) == 4, f"Bounding box must have 4 elements, got f{len(bbox)}"
+        x, y, w, h = bbox
+        return BBox(title=title or "", x1=x, y1=y, x2=x + w, y2=y + h)

datachain/lib/models/pose.py

@@ -8,48 +8,15 @@ 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.
+        x (list[float]): The x-coordinates of the keypoints.
+        y (list[float]): 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 coordinates must be a list of 2 lists."
-        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 integers or floats."
-        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 set(points) == {
-            "x",
-            "y",
-        }, "Pose coordinates must contain keys 'x' and 'y'."
-        points_x, points_y = points["x"], points["y"]
-        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 integers or floats."
-        return Pose(
-            x=[round(coord) for coord in points_x],
-            y=[round(coord) for coord in points_y],
-        )
+    x: list[float] = Field(default=None)
+    y: list[float] = Field(default=None)
 
 
 class Pose3D(DataModel):
@@ -57,52 +24,14 @@ 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.
+        x (list[float]): The x-coordinates of the keypoints.
+        y (list[float]): 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)
+    x: list[float] = Field(default=None)
+    y: list[float] = Field(default=None)
     visible: list[float] = Field(default=None)
-
-    @staticmethod
-    def from_list(points: list[list[float]]) -> "Pose3D":
-        assert len(points) == 3, "Pose coordinates must be a list of 3 lists."
-        points_x, points_y, points_v = points
-        assert (
-            len(points_x) == len(points_y) == len(points_v) == 17
-        ), "Pose x, y, and visibility coordinates must have the same length of 17."
-        assert all(
-            isinstance(value, (int, float))
-            for value in [*points_x, *points_y, *points_v]
-        ), "Pose coordinates must be integers or floats."
-        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 set(points) == {
-            "x",
-            "y",
-            "visible",
-        }, "Pose coordinates must contain keys 'x', 'y', and 'visible'."
-        points_x, points_y, points_v = points["x"], points["y"], points["visible"]
-        assert (
-            len(points_x) == len(points_y) == len(points_v) == 17
-        ), "Pose x, y, and visibility coordinates must have the same length of 17."
-        assert all(
-            isinstance(value, (int, float))
-            for value in [*points_x, *points_y, *points_v]
-        ), "Pose coordinates must be integers or floats."
-        return Pose3D(
-            x=[round(coord) for coord in points_x],
-            y=[round(coord) for coord in points_y],
-            visible=points_v,
-        )

datachain/lib/models/yolo.py

@@ -0,0 +1,39 @@
+"""
+This module contains the YOLO models.
+
+YOLO stands for "You Only Look Once", a family of object detection models that
+are designed to be fast and accurate. The models are trained to detect objects
+in images by dividing the image into a grid and predicting the bounding boxes
+and class probabilities for each grid cell.
+
+More information about YOLO can be found here:
+- https://pjreddie.com/darknet/yolo/
+- https://docs.ultralytics.com/
+"""
+
+
+class PoseBodyPart:
+    """
+    An enumeration of body parts for YOLO pose keypoints.
+
+    More information about the body parts can be found here:
+    https://docs.ultralytics.com/tasks/pose/
+    """
+
+    nose = 0
+    left_eye = 1
+    right_eye = 2
+    left_ear = 3
+    right_ear = 4
+    left_shoulder = 5
+    right_shoulder = 6
+    left_elbow = 7
+    right_elbow = 8
+    left_wrist = 9
+    right_wrist = 10
+    left_hip = 11
+    right_hip = 12
+    left_knee = 13
+    right_knee = 14
+    left_ankle = 15
+    right_ankle = 16

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/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/node.py

@@ -55,7 +55,7 @@ class Node:
     last_modified: Optional[datetime] = None
     size: int = 0
     location: Optional[str] = None
-    source: StorageURI = StorageURI("")
+    source: StorageURI = StorageURI("")  # noqa: RUF009
     dir_type: int = DirType.FILE
 
     @property

datachain/query/dataset.py

@@ -473,33 +473,31 @@ class UDFStep(Step, ABC):
                 # Otherwise process single-threaded (faster for smaller UDFs)
                 warehouse = self.catalog.warehouse
 
-                with contextlib.closing(
-                    batching(warehouse.dataset_select_paginated, query)
-                ) as udf_inputs:
-                    download_cb = get_download_callback()
-                    processed_cb = get_processed_callback()
-                    generated_cb = get_generated_callback(self.is_generator)
-                    try:
-                        udf_results = self.udf.run(
-                            udf_fields,
-                            udf_inputs,
-                            self.catalog,
-                            self.is_generator,
-                            self.cache,
-                            download_cb,
-                            processed_cb,
-                        )
-                        process_udf_outputs(
-                            warehouse,
-                            udf_table,
-                            udf_results,
-                            self.udf,
-                            cb=generated_cb,
-                        )
-                    finally:
-                        download_cb.close()
-                        processed_cb.close()
-                        generated_cb.close()
+                udf_inputs = batching(warehouse.dataset_select_paginated, query)
+                download_cb = get_download_callback()
+                processed_cb = get_processed_callback()
+                generated_cb = get_generated_callback(self.is_generator)
+                try:
+                    udf_results = self.udf.run(
+                        udf_fields,
+                        udf_inputs,
+                        self.catalog,
+                        self.is_generator,
+                        self.cache,
+                        download_cb,
+                        processed_cb,
+                    )
+                    process_udf_outputs(
+                        warehouse,
+                        udf_table,
+                        udf_results,
+                        self.udf,
+                        cb=generated_cb,
+                    )
+                finally:
+                    download_cb.close()
+                    processed_cb.close()
+                    generated_cb.close()
 
                 warehouse.insert_rows_done(udf_table)
 

datachain/sql/sqlite/types.py

@@ -36,7 +36,14 @@ def convert_array(arr):
 
 
 def adapt_np_array(arr):
-    return orjson.dumps(arr, option=orjson.OPT_SERIALIZE_NUMPY).decode("utf-8")
+    def _json_serialize(obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        return obj
+
+    return orjson.dumps(
+        arr, option=orjson.OPT_SERIALIZE_NUMPY, default=_json_serialize
+    ).decode("utf-8")
 
 
 def adapt_np_generic(val):

{datachain-0.6.10.dist-info → datachain-0.7.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: datachain
-Version: 0.6.10
+Version: 0.7.0
 Summary: Wrangle unstructured AI data at scale
 Author-email: Dmitry Petrov <support@dvc.org>
 License: Apache-2.0

{datachain-0.6.10.dist-info → datachain-0.7.0.dist-info}/RECORD

@@ -1,6 +1,6 @@
 datachain/__init__.py,sha256=nnTyB5MpCfBZ6D85JPz-5hUT7i-68Is-47Bxgew8lRw,930
 datachain/__main__.py,sha256=hG3Y4ARGEqe1AWwNMd259rBlqtphx1Wk39YbueQ0yV8,91
-datachain/asyn.py,sha256=Lg3Ck1PQLjQziMx9KU4atzbEnJXTE0924WMYkhgWtGU,8247
+datachain/asyn.py,sha256=5aKrjnUxk0mtnZeFKNJd1DCE0MsnSoyJBZkr0y9H_a0,9313
 datachain/cache.py,sha256=s0YHN7qurmQv-eC265TjeureK84TebWWAnL07cxchZQ,2997
 datachain/cli.py,sha256=hdVt_HJumQVgtaBAtBVJm-uPyYVogMXNVLmRcZyWHgk,36677
 datachain/cli_utils.py,sha256=jrn9ejGXjybeO1ur3fjdSiAyCHZrX0qsLLbJzN9ErPM,2418
@@ -9,7 +9,7 @@ datachain/dataset.py,sha256=0IN-5y723y-bnFlieKtOFZLCjwX_yplFo3q0DV7LRPw,14821
 datachain/error.py,sha256=bxAAL32lSeMgzsQDEHbGTGORj-mPzzpCRvWDPueJNN4,1092
 datachain/job.py,sha256=Jt4sNutMHJReaGsj3r3scueN5aESLGfhimAa8pUP7Is,1271
 datachain/listing.py,sha256=TgKg25ZWAP5enzKgw2_2GUPJVdnQUh6uySHB5SJrUY4,7773
-datachain/node.py,sha256=i7_jC8VcW6W5VYkDszAOu0H-rNBuqXB4UnLEh4wFzjc,5195
+datachain/node.py,sha256=o8Sqy92QkzzcLK6XmIFLyDSE6Rw6kUTmGRhEmfLFdhg,5211
 datachain/nodes_fetcher.py,sha256=F-73-h19HHNGtHFBGKk7p3mc0ALm4a9zGnzhtuUjnp4,1107
 datachain/nodes_thread_pool.py,sha256=uPo-xl8zG5m9YgODjPFBpbcqqHjI-dcxH87yAbj_qco,3192
 datachain/progress.py,sha256=5KotcvvzAUL_RF0GEj4JY0IB1lyImnmHxe89YkT1XO4,4330
@@ -37,14 +37,14 @@ datachain/data_storage/metastore.py,sha256=5b7o_CSHC2djottebYn-Hq5q0yaSLOKPIRCna
 datachain/data_storage/schema.py,sha256=scANMQqozita3HjEtq7eupMgh6yYkrZHoXtfuL2RoQg,9879
 datachain/data_storage/serializer.py,sha256=6G2YtOFqqDzJf1KbvZraKGXl2XHZyVml2krunWUum5o,927
 datachain/data_storage/sqlite.py,sha256=CspRUlYsIcubgzvcQxTACnmcuKESSLZcqCl0dcrtRiA,27471
-datachain/data_storage/warehouse.py,sha256=xwMaR4jBpR13vjG3zrhphH4z2_CFLNj0KPF0LJCXCJ8,30727
+datachain/data_storage/warehouse.py,sha256=yXNU0U3exzR1E6dqbYYmL4RhXWsbYWVdZ3jONGcVniY,30914
 datachain/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datachain/lib/arrow.py,sha256=-hu9tic79a01SY2UBqkA3U6wUr6tnE3T3q5q_BnO93A,9156
 datachain/lib/clip.py,sha256=lm5CzVi4Cj1jVLEKvERKArb-egb9j1Ls-fwTItT6vlI,6150
 datachain/lib/data_model.py,sha256=dau4AlZBhOFvF7pEKMeqCeRkcFFg5KFvTBWW_2CdH5g,2371
 datachain/lib/dataset_info.py,sha256=q0EW9tj5jXGSD9Lzct9zbH4P1lfIGd_cIWqhnMxv7Q0,2464
-datachain/lib/dc.py,sha256=BmRgCt5fXvBqlFV07KN-nWszueRyCkC7td1x7T4BZ7k,87688
-datachain/lib/file.py,sha256=lHxE1wOGR4QJBQ3AYjhPLwpX72dOi06vkcwA-WSAGlg,14817
+datachain/lib/dc.py,sha256=u0RQJPG0zwxsoYS-4wrbDBPuLYZajwIi1YX37khKfkI,87942
+datachain/lib/file.py,sha256=-XMkL6ED1sE7TMhWoMRTEuOXswZJw8X6AEmJDONFP74,15019
 datachain/lib/hf.py,sha256=BW2NPpqxkpPwkSaGlppT8Rbs8zPpyYC-tR6htY08c-0,5817
 datachain/lib/image.py,sha256=AMXYwQsmarZjRbPCZY3M1jDsM2WAB_b3cTY4uOIuXNU,2675
 datachain/lib/listing.py,sha256=cVkCp7TRVpcZKSx-Bbk9t51bQI9Mw0o86W6ZPhAsuzM,3667
@@ -52,11 +52,11 @@ datachain/lib/listing_info.py,sha256=9ua40Hw0aiQByUw3oAEeNzMavJYfW0Uhe8YdCTK-m_g
 datachain/lib/meta_formats.py,sha256=anK2bDVbaeCCh0yvKUBaW2MVos3zRgdaSV8uSduzPcU,6680
 datachain/lib/model_store.py,sha256=DNIv8Y6Jtk1_idNLzIpsThOsdW2BMAudyUCbPUcgcxk,2515
 datachain/lib/pytorch.py,sha256=W-ARi2xH1f1DUkVfRuerW-YWYgSaJASmNCxtz2lrJGI,6072
-datachain/lib/settings.py,sha256=39thOpYJw-zPirzeNO6pmRC2vPrQvt4eBsw1xLWDFsw,2344
+datachain/lib/settings.py,sha256=ZELRCTLbi5vzRPiDX6cQ9LLg9TefJ_A05gIGni0lll8,2535
 datachain/lib/signal_schema.py,sha256=xwkE5bxJxUhZTjrA6jqN87XbSXPikCbL6eOPL9WyrKM,24556
 datachain/lib/tar.py,sha256=3WIzao6yD5fbLqXLTt9GhPGNonbFIs_fDRu-9vgLgsA,1038
 datachain/lib/text.py,sha256=UNHm8fhidk7wdrWqacEWaA6I9ykfYqarQ2URby7jc7M,1261
-datachain/lib/udf.py,sha256=4CqK51n3bntXCmkwoOQIrX34wMKOknkC23HtR4D_2vM,12705
+datachain/lib/udf.py,sha256=-j0krjNAELTqRI0dB1N65AmawtcIY5vN---AuUcW8Us,13637
 datachain/lib/udf_signature.py,sha256=GXw24A-Olna6DWCdgy2bC-gZh_gLGPQ-KvjuI6pUjC0,7281
 datachain/lib/utils.py,sha256=6NwgWLl5JrgtD4rsSFEe-yR2ntEwJMJEtAZ3FIxK3fg,1529
 datachain/lib/vfile.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -71,17 +71,13 @@ datachain/lib/convert/values_to_tuples.py,sha256=varRCnSMT_pZmHznrd2Yi05qXLLz_v9
 datachain/lib/func/__init__.py,sha256=wlAKhGV0QDg9y7reSwoUF8Vicfqh_YOUNIXLzxICGz4,403
 datachain/lib/func/aggregate.py,sha256=H1ziFQdaK9zvnxvttfnEzkkyGvEEmMAvmgCsBV6nfm8,10917
 datachain/lib/func/func.py,sha256=HAJZ_tpiRG2R-et7pr0WnoyNZYtpbPn3_HBuL3RQpbU,4800
-datachain/lib/models/__init__.py,sha256=6iwqXWcybyELKdLEe59yUPl8R8ZHDY4lA-xCHVYPdOA,191
-datachain/lib/models/bbox.py,sha256=UJ_64D8TQglX2B_ueseILPoT3cGIWr9McVg0mv2YdmE,3717
-datachain/lib/models/pose.py,sha256=KC-OpLC7-3v6qg4YN6pXlfAgtg88VLQoRc75JCEmbfY,3931
-datachain/lib/models/segment.py,sha256=ergCFnEzLDzaU75p1_KvWgal1LSv4VuFmkWLkRJeaVk,1862
-datachain/lib/models/ultralytics/__init__.py,sha256=g8mgII0k_RJiOG9kd4k_ECfCgDhT_iPh3vCC_5OiDD4,305
-datachain/lib/models/ultralytics/bbox.py,sha256=LAaezAnnugfBiczWZ63NTo65kX2BegR5WGXjQTOTE28,5784
-datachain/lib/models/ultralytics/pose.py,sha256=nMoEeeY_Zi7Iiu7vIo9ZTq8ARUdg_BcZMQIA_WgRNk4,3488
-datachain/lib/models/ultralytics/segment.py,sha256=IHnthsq6uQ6DSdHLK2akbdd0Eq8wW7oaAK6pUG8nxJc,3818
+datachain/lib/models/__init__.py,sha256=AGvjPbUokJiir3uelTa4XGtNSECkMFc5Xmi_N3AtxPQ,119
+datachain/lib/models/bbox.py,sha256=aiYNhvEcRK3dEN4MBcptmkPKc9kMP16ZQdu7xPk6hek,1555
+datachain/lib/models/pose.py,sha256=peuJPNSiGuTXfCfGIABwv8PGYistvTTBmtf-8X8E_eA,1077
+datachain/lib/models/yolo.py,sha256=eftoJDUa8iOpFTF1EkKVAd5Q-3HRd6X4eCIZ9h5p4nI,972
 datachain/query/__init__.py,sha256=7DhEIjAA8uZJfejruAVMZVcGFmvUpffuZJwgRqNwe-c,263
 datachain/query/batch.py,sha256=5fEhORFe7li12SdYddaSK3LyqksMfCHhwN1_A6TfsA4,3485
-datachain/query/dataset.py,sha256=MGArYxioeGvm8w7hQtQAjEI6wsZN_XAoh4-jO4d0U5Q,53926
+datachain/query/dataset.py,sha256=sQny-ZemB2HueC4mPg-7qSaqUD85MMO-DQyVVP8K1CA,53765
 datachain/query/dispatch.py,sha256=wjjTWw6sFQbB9SKRh78VbfvwSMgJXCfqJklS3-9KnCU,12025
 datachain/query/metrics.py,sha256=r5b0ygYhokbXp8Mg3kCH8iFSRw0jxzyeBe-C-J_bKFc,938
 datachain/query/params.py,sha256=O_j89mjYRLOwWNhYZl-z7mi-rkdP7WyFmaDufsdTryE,863
@@ -105,14 +101,14 @@ datachain/sql/functions/random.py,sha256=vBwEEj98VH4LjWixUCygQ5Bz1mv1nohsCG0-ZTE
 datachain/sql/functions/string.py,sha256=DYgiw8XSk7ge7GXvyRI1zbaMruIizNeI-puOjriQGZQ,1148
 datachain/sql/sqlite/__init__.py,sha256=TAdJX0Bg28XdqPO-QwUVKy8rg78cgMileHvMNot7d04,166
 datachain/sql/sqlite/base.py,sha256=aHSZVvh4XSVkvZ07h3jMoRlHI4sWD8y3SnmGs9xMG9Y,14375
-datachain/sql/sqlite/types.py,sha256=yzvp0sXSEoEYXs6zaYC_2YubarQoZH-MiUNXcpuEP4s,1573
+datachain/sql/sqlite/types.py,sha256=lPXS1XbkmUtlkkiRxy_A_UzsgpPv2VSkXYOD4zIHM4w,1734
 datachain/sql/sqlite/vector.py,sha256=ncW4eu2FlJhrP_CIpsvtkUabZlQdl2D5Lgwy_cbfqR0,469
 datachain/toolkit/__init__.py,sha256=eQ58Q5Yf_Fgv1ZG0IO5dpB4jmP90rk8YxUWmPc1M2Bo,68
 datachain/toolkit/split.py,sha256=6FcEJgUsJsUcCqKW5aXuJy4DvbcQ7_dFbsfNPhn8EVg,2377
 datachain/torch/__init__.py,sha256=gIS74PoEPy4TB3X6vx9nLO0Y3sLJzsA8ckn8pRWihJM,579
-datachain-0.6.10.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
-datachain-0.6.10.dist-info/METADATA,sha256=AgQuuefAhZRIL1jDJWz-q4daqA5ZmnQN8dafqnt01XA,18038
-datachain-0.6.10.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
-datachain-0.6.10.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
-datachain-0.6.10.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
-datachain-0.6.10.dist-info/RECORD,,
+datachain-0.7.0.dist-info/LICENSE,sha256=8DnqK5yoPI_E50bEg_zsHKZHY2HqPy4rYN338BHQaRA,11344
+datachain-0.7.0.dist-info/METADATA,sha256=Cm0v22C-aT14JBG4NjNQmM7kvhYZFAf6rh0yDAMRWFU,18037
+datachain-0.7.0.dist-info/WHEEL,sha256=R06PA3UVYHThwHvxuRWMqaGcr-PuniXahwjmQRFMEkY,91
+datachain-0.7.0.dist-info/entry_points.txt,sha256=0GMJS6B_KWq0m3VT98vQI2YZodAMkn4uReZ_okga9R4,49
+datachain-0.7.0.dist-info/top_level.txt,sha256=lZPpdU_2jJABLNIg2kvEOBi8PtsYikbN1OdMLHk8bTg,10
+datachain-0.7.0.dist-info/RECORD,,

datachain/lib/models/segment.py

@@ -1,53 +0,0 @@
-from pydantic import Field
-
-from datachain.lib.data_model import DataModel
-
-
-class Segments(DataModel):
-    """
-    A data model for representing segments.
-
-    Attributes:
-        title (str): The title of the segments.
-        x (list[int]): The x-coordinates of the segments.
-        y (list[int]): The y-coordinates of the segments.
-
-    The segments are represented as lists of x and y coordinates, where each index
-    corresponds to a specific segment.
-    """
-
-    title: str = Field(default="")
-    x: list[int] = Field(default=None)
-    y: list[int] = Field(default=None)
-
-    @staticmethod
-    def from_list(points: list[list[float]], title: str = "") -> "Segments":
-        assert len(points) == 2, "Segments coordinates must be a list of 2 lists."
-        points_x, points_y = points
-        assert len(points_x) == len(
-            points_y
-        ), "Segments x and y coordinates must have the same length."
-        assert all(
-            isinstance(value, (int, float)) for value in [*points_x, *points_y]
-        ), "Segments coordinates must be integers or floats."
-        return Segments(
-            title=title,
-            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]], title: str = "") -> "Segments":
-        assert set(points) == {
-            "x",
-            "y",
-        }, "Segments coordinates must contain keys 'x' and 'y'."
-        points_x, points_y = points["x"], points["y"]
-        assert all(
-            isinstance(value, (int, float)) for value in [*points_x, *points_y]
-        ), "Segments coordinates must be integers or floats."
-        return Segments(
-            title=title,
-            x=[round(coord) for coord in points_x],
-            y=[round(coord) for coord in points_y],
-        )

datachain/lib/models/ultralytics/__init__.py

@@ -1,14 +0,0 @@
-from .bbox import YoloBBox, YoloBBoxes, YoloOBBox, YoloOBBoxes
-from .pose import YoloPose, YoloPoses
-from .segment import YoloSegment, YoloSegments
-
-__all__ = [
-    "YoloBBox",
-    "YoloBBoxes",
-    "YoloOBBox",
-    "YoloOBBoxes",
-    "YoloPose",
-    "YoloPoses",
-    "YoloSegment",
-    "YoloSegments",
-]

datachain/lib/models/ultralytics/bbox.py

@@ -1,189 +0,0 @@
-"""
-This module contains the YOLO models.
-
-YOLO stands for "You Only Look Once", a family of object detection models that
-are designed to be fast and accurate. The models are trained to detect objects
-in images by dividing the image into a grid and predicting the bounding boxes
-and class probabilities for each grid cell.
-
-More information about YOLO can be found here:
-- https://pjreddie.com/darknet/yolo/
-- https://docs.ultralytics.com/
-"""
-
-from io import BytesIO
-from typing import TYPE_CHECKING
-
-from PIL import Image
-from pydantic import Field
-
-from datachain.lib.data_model import DataModel
-from datachain.lib.models.bbox import BBox, OBBox
-
-if TYPE_CHECKING:
-    from ultralytics.engine.results import Results
-    from ultralytics.models import YOLO
-
-    from datachain.lib.file import File
-
-
-class YoloBBox(DataModel):
-    """
-    A class representing a bounding box detected by a YOLO model.
-
-    Attributes:
-        cls: The class of the detected object.
-        name: The name of the detected object.
-        confidence: The confidence score of the detection.
-        box: The bounding box of the detected object
-    """
-
-    cls: int = Field(default=-1)
-    name: str = Field(default="")
-    confidence: float = Field(default=0)
-    box: BBox = Field(default=None)
-
-    @staticmethod
-    def from_file(yolo: "YOLO", file: "File") -> "YoloBBox":
-        results = yolo(Image.open(BytesIO(file.read())))
-        if len(results) == 0:
-            return YoloBBox()
-        return YoloBBox.from_result(results[0])
-
-    @staticmethod
-    def from_result(result: "Results") -> "YoloBBox":
-        summary = result.summary()
-        if not summary:
-            return YoloBBox()
-        name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
-            else BBox()
-        )
-        return YoloBBox(
-            cls=summary[0]["class"],
-            name=name,
-            confidence=summary[0]["confidence"],
-            box=box,
-        )
-
-
-class YoloBBoxes(DataModel):
-    """
-    A class representing a list of bounding boxes detected by a YOLO model.
-
-    Attributes:
-        cls: A list of classes of the detected objects.
-        name: A list of names of the detected objects.
-        confidence: A list of confidence scores of the detections.
-        box: A list of bounding boxes of the detected objects
-    """
-
-    cls: list[int]
-    name: list[str]
-    confidence: list[float]
-    box: list[BBox]
-
-    @staticmethod
-    def from_file(yolo: "YOLO", file: "File") -> "YoloBBoxes":
-        results = yolo(Image.open(BytesIO(file.read())))
-        return YoloBBoxes.from_results(results)
-
-    @staticmethod
-    def from_results(results: list["Results"]) -> "YoloBBoxes":
-        cls, names, confidence, box = [], [], [], []
-        for r in results:
-            for s in r.summary():
-                name = s.get("name", "")
-                cls.append(s["class"])
-                names.append(name)
-                confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
-        return YoloBBoxes(
-            cls=cls,
-            name=names,
-            confidence=confidence,
-            box=box,
-        )
-
-
-class YoloOBBox(DataModel):
-    """
-    A class representing an oriented bounding box detected by a YOLO model.
-
-    Attributes:
-        cls: The class of the detected object.
-        name: The name of the detected object.
-        confidence: The confidence score of the detection.
-        box: The oriented bounding box of the detected object.
-    """
-
-    cls: int = Field(default=-1)
-    name: str = Field(default="")
-    confidence: float = Field(default=0)
-    box: OBBox = Field(default=None)
-
-    @staticmethod
-    def from_file(yolo: "YOLO", file: "File") -> "YoloOBBox":
-        results = yolo(Image.open(BytesIO(file.read())))
-        if len(results) == 0:
-            return YoloOBBox()
-        return YoloOBBox.from_result(results[0])
-
-    @staticmethod
-    def from_result(result: "Results") -> "YoloOBBox":
-        summary = result.summary()
-        if not summary:
-            return YoloOBBox()
-        name = summary[0].get("name", "")
-        box = (
-            OBBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
-            else OBBox()
-        )
-        return YoloOBBox(
-            cls=summary[0]["class"],
-            name=name,
-            confidence=summary[0]["confidence"],
-            box=box,
-        )
-
-
-class YoloOBBoxes(DataModel):
-    """
-    A class representing a list of oriented bounding boxes detected by a YOLO model.
-
-    Attributes:
-        cls: A list of classes of the detected objects.
-        name: A list of names of the detected objects.
-        confidence: A list of confidence scores of the detections.
-        box: A list of oriented bounding boxes of the detected objects.
-    """
-
-    cls: list[int]
-    name: list[str]
-    confidence: list[float]
-    box: list[OBBox]
-
-    @staticmethod
-    def from_file(yolo: "YOLO", file: "File") -> "YoloOBBoxes":
-        results = yolo(Image.open(BytesIO(file.read())))
-        return YoloOBBoxes.from_results(results)
-
-    @staticmethod
-    def from_results(results: list["Results"]) -> "YoloOBBoxes":
-        cls, names, confidence, box = [], [], [], []
-        for r in results:
-            for s in r.summary():
-                name = s.get("name", "")
-                cls.append(s["class"])
-                names.append(name)
-                confidence.append(s["confidence"])
-                box.append(OBBox.from_dict(s.get("box", {}), title=name))
-        return YoloOBBoxes(
-            cls=cls,
-            name=names,
-            confidence=confidence,
-            box=box,
-        )

datachain/lib/models/ultralytics/pose.py

@@ -1,126 +0,0 @@
-"""
-This module contains the YOLO models.
-
-YOLO stands for "You Only Look Once", a family of object detection models that
-are designed to be fast and accurate. The models are trained to detect objects
-in images by dividing the image into a grid and predicting the bounding boxes
-and class probabilities for each grid cell.
-
-More information about YOLO can be found here:
-- https://pjreddie.com/darknet/yolo/
-- https://docs.ultralytics.com/
-"""
-
-from typing import TYPE_CHECKING
-
-from pydantic import Field
-
-from datachain.lib.data_model import DataModel
-from datachain.lib.models.bbox import BBox
-from datachain.lib.models.pose import Pose3D
-
-if TYPE_CHECKING:
-    from ultralytics.engine.results import Results
-
-
-class YoloPoseBodyPart:
-    """An enumeration of body parts for YOLO pose keypoints."""
-
-    nose = 0
-    left_eye = 1
-    right_eye = 2
-    left_ear = 3
-    right_ear = 4
-    left_shoulder = 5
-    right_shoulder = 6
-    left_elbow = 7
-    right_elbow = 8
-    left_wrist = 9
-    right_wrist = 10
-    left_hip = 11
-    right_hip = 12
-    left_knee = 13
-    right_knee = 14
-    left_ankle = 15
-    right_ankle = 16
-
-
-class YoloPose(DataModel):
-    """
-    A data model for YOLO pose keypoints.
-
-    Attributes:
-        cls: The class of the pose.
-        name: The name of the pose.
-        confidence: The confidence score of the pose.
-        box: The bounding box of the pose.
-        keypoints: The 3D pose keypoints.
-    """
-
-    cls: int = Field(default=-1)
-    name: str = Field(default="")
-    confidence: float = Field(default=0)
-    box: BBox = Field(default=None)
-    keypoints: Pose3D = Field(default=None)
-
-    @staticmethod
-    def from_result(result: "Results") -> "YoloPose":
-        summary = result.summary()
-        if not summary:
-            return YoloPose()
-        name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
-            else BBox()
-        )
-        keypoints = (
-            Pose3D.from_dict(summary[0]["keypoints"])
-            if "keypoints" in summary[0]
-            else Pose3D()
-        )
-        return YoloPose(
-            cls=summary[0]["class"],
-            name=name,
-            confidence=summary[0]["confidence"],
-            box=box,
-            keypoints=keypoints,
-        )
-
-
-class YoloPoses(DataModel):
-    """
-    A data model for a list of YOLO pose keypoints.
-
-    Attributes:
-        cls: The classes of the poses.
-        name: The names of the poses.
-        confidence: The confidence scores of the poses.
-        box: The bounding boxes of the poses.
-        keypoints: The 3D pose keypoints of the poses.
-    """
-
-    cls: list[int]
-    name: list[str]
-    confidence: list[float]
-    box: list[BBox]
-    keypoints: list[Pose3D]
-
-    @staticmethod
-    def from_results(results: list["Results"]) -> "YoloPoses":
-        cls, names, confidence, box, keypoints = [], [], [], [], []
-        for r in results:
-            for s in r.summary():
-                name = s.get("name", "")
-                cls.append(s["class"])
-                names.append(name)
-                confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
-                keypoints.append(Pose3D.from_dict(s.get("keypoints", {})))
-        return YoloPoses(
-            cls=cls,
-            name=names,
-            confidence=confidence,
-            box=box,
-            keypoints=keypoints,
-        )

datachain/lib/models/ultralytics/segment.py

@@ -1,121 +0,0 @@
-"""
-This module contains the YOLO models.
-
-YOLO stands for "You Only Look Once", a family of object detection models that
-are designed to be fast and accurate. The models are trained to detect objects
-in images by dividing the image into a grid and predicting the bounding boxes
-and class probabilities for each grid cell.
-
-More information about YOLO can be found here:
-- https://pjreddie.com/darknet/yolo/
-- https://docs.ultralytics.com/
-"""
-
-from io import BytesIO
-from typing import TYPE_CHECKING
-
-from PIL import Image
-from pydantic import Field
-
-from datachain.lib.data_model import DataModel
-from datachain.lib.models.bbox import BBox
-from datachain.lib.models.segment import Segments
-
-if TYPE_CHECKING:
-    from ultralytics.engine.results import Results
-    from ultralytics.models import YOLO
-
-    from datachain.lib.file import File
-
-
-class YoloSegment(DataModel):
-    """
-    A data model for a single YOLO segment.
-
-    Attributes:
-        cls (int): The class of the segment.
-        name (str): The name of the segment.
-        confidence (float): The confidence of the segment.
-        box (BBox): The bounding box of the segment.
-        segments (Segments): The segments of the segment.
-    """
-
-    cls: int = Field(default=-1)
-    name: str = Field(default="")
-    confidence: float = Field(default=0)
-    box: BBox = Field(default=None)
-    segments: Segments = Field(default=None)
-
-    @staticmethod
-    def from_file(yolo: "YOLO", file: "File") -> "YoloSegment":
-        results = yolo(Image.open(BytesIO(file.read())))
-        if len(results) == 0:
-            return YoloSegment()
-        return YoloSegment.from_result(results[0])
-
-    @staticmethod
-    def from_result(result: "Results") -> "YoloSegment":
-        summary = result.summary()
-        if not summary:
-            return YoloSegment()
-        name = summary[0].get("name", "")
-        box = (
-            BBox.from_dict(summary[0]["box"], title=name)
-            if "box" in summary[0]
-            else BBox()
-        )
-        segments = (
-            Segments.from_dict(summary[0]["segments"], title=name)
-            if "segments" in summary[0]
-            else Segments()
-        )
-        return YoloSegment(
-            cls=summary[0]["class"],
-            name=summary[0]["name"],
-            confidence=summary[0]["confidence"],
-            box=box,
-            segments=segments,
-        )
-
-
-class YoloSegments(DataModel):
-    """
-    A data model for a list of YOLO segments.
-
-    Attributes:
-        cls (list[int]): The classes of the segments.
-        name (list[str]): The names of the segments.
-        confidence (list[float]): The confidences of the segments.
-        box (list[BBox]): The bounding boxes of the segments.
-        segments (list[Segments]): The segments of the segments.
-    """
-
-    cls: list[int]
-    name: list[str]
-    confidence: list[float]
-    box: list[BBox]
-    segments: list[Segments]
-
-    @staticmethod
-    def from_file(yolo: "YOLO", file: "File") -> "YoloSegments":
-        results = yolo(Image.open(BytesIO(file.read())))
-        return YoloSegments.from_results(results)
-
-    @staticmethod
-    def from_results(results: list["Results"]) -> "YoloSegments":
-        cls, names, confidence, box, segments = [], [], [], [], []
-        for r in results:
-            for s in r.summary():
-                name = s.get("name", "")
-                cls.append(s["class"])
-                names.append(name)
-                confidence.append(s["confidence"])
-                box.append(BBox.from_dict(s.get("box", {}), title=name))
-                segments.append(Segments.from_dict(s.get("segments", {}), title=name))
-        return YoloSegments(
-            cls=cls,
-            name=names,
-            confidence=confidence,
-            box=box,
-            segments=segments,
-        )