Version 0.0.1-alpha.0.

Author: ATATC

leads_jarvis/checkpoints/yolov8m.pt

@@ -1,8 +1,11 @@
 from abc import ABCMeta as _ABCMeta, abstractmethod as _abstractmethod
 from typing import Any as _Any
 
+from cv2 import rectangle as _rectangle, putText as _put_text, FONT_HERSHEY_COMPLEX_SMALL as _FONT
 from numpy import ndarray as _ndarray
 
+from leads_jarvis.utils import to_opencv
+
 
 class Detection(object, metaclass=_ABCMeta):
     @_abstractmethod
@@ -11,3 +14,18 @@ def detect(self, image: _ndarray) -> list[dict[str, _Any]]:
 
     def __call__(self, image: _ndarray) -> list[dict[str, _Any]]:
         return self.detect(image)
+
+    def mark(self, image: _ndarray, data: list[dict[str, _Any]] | None = None,
+             filter_type: tuple[str, ...] | None = None) -> _ndarray:
+        if data is None:
+            data = self(image)
+        image = to_opencv(image)
+        for item in data:
+            name = item["name"]
+            if filter_type is not None and name not in filter_type:
+                continue
+            box = item["box"]
+            x1, y1, x2, y2 = round(box["x1"]), round(box["y1"]), round(box["x2"]), round(box["y2"])
+            image = _rectangle(image, (x1, y1), (x2, y2), (255, 255, 255))
+            image = _put_text(image, name, (x1, y1 - 4), _FONT, 1, (255, 255, 255))
+        return image

leads_jarvis/checkpoints/yolov8n.pt

@@ -4,7 +4,7 @@
 from ultralytics import YOLO as _YOLO
 
 from leads_jarvis.network.detection.prototype import Detection
-from leads_jarvis.utils import _CHECKPOINTS_PATH
+from leads_jarvis.utils import _CHECKPOINTS_PATH, to_opencv
 
 
 class PretrainedYOLOWrapper(Detection):
@@ -13,7 +13,7 @@ def __init__(self, yolo: _YOLO) -> None:
 
     @_override
     def detect(self, image: _ndarray) -> list[dict[str, _Any]]:
-        return self._yolo(image.transpose(1, 2, 0))[0].summary()
+        return self._yolo(to_opencv(image))[0].summary()
 
 
 class PretrainedYOLO(Detection):

leads_jarvis/checkpoints/yolov8s.pt

@@ -5,7 +5,7 @@
     SamAutomaticMaskGenerator as _SamAutomaticMaskGenerator
 
 from leads_jarvis.network.segmentation.prototype import Segmentation
-from leads_jarvis.utils import download_checkpoint
+from leads_jarvis.utils import download_checkpoint, to_opencv
 
 
 class PretrainedSAMWrapper(Segmentation):
@@ -15,7 +15,7 @@ def __init__(self, mask_generator: _SamAutomaticMaskGenerator) -> None:
 
     @_override
     def segment(self, image: _ndarray) -> list[dict[str, _Any]]:
-        return self._mask_generator.generate(image.transpose(1, 2, 0))
+        return self._mask_generator.generate(to_opencv(image))
 
 
 class PretrainedSAM(Segmentation):

leads_jarvis/network/detection/prototype.py

@@ -1,12 +1,21 @@
 from os.path import abspath as _abspath, exists as _exists
 
 from leads import L as _L
+from numpy import ndarray as _ndarray
 from requests import head as _head, get as _get
 from rich.progress import Progress as _Progress
 
 _CHECKPOINTS_PATH: str = f"{_abspath(__file__)[:-8]}checkpoints"
 
 
+def from_opencv(image: _ndarray) -> _ndarray:
+    return image.transpose(2, 0, 1)
+
+
+def to_opencv(image: _ndarray) -> _ndarray:
+    return image.transpose(1, 2, 0)
+
+
 def download_checkpoint(url: str, to: str, overwrite: bool = False) -> str:
     if _exists(to := f"{_CHECKPOINTS_PATH}/{to}") and not overwrite:
         return to

leads_jarvis/network/detection/yolo.py

@@ -5,7 +5,7 @@
 
 setup(
     name="leads-jarvis",
-    version="0.0.1",
+    version="0.0.1-alpha.0",
     python_requires=">=3.12",
     author="ProjectNeura",
     author_email="central@projectneura.org",
@@ -16,8 +16,9 @@
     url="https://github.com/ProjectNeura/LEADS-Jarvis",
     packages=find_packages(),
     package_data={
-        "leads_jarvis": ["weights/*"]
+        "leads_jarvis": ["checkpoints/*"]
     },
     include_package_data=True,
-    install_requires=["leads>=0.9.1-alpha.3", "torch", "torchvision", "timm"]
+    install_requires=["leads>=0.9.1", "torch", "torchvision", "timm", "segment-anything", "requests", "rich",
+                      "ultralytics"]
 )