From 2da9f913f3b4d05f6ce02819751fc74feef50796 Mon Sep 17 00:00:00 2001 From: Ryan Dick Date: Tue, 30 Jul 2024 16:04:29 -0400 Subject: [PATCH] Add detection_result.py - was forgotten in a prior commit --- .../backend/grounded_sam/detection_result.py | 41 +++++++++++++++++++ 1 file changed, 41 insertions(+) create mode 100644 invokeai/backend/grounded_sam/detection_result.py diff --git a/invokeai/backend/grounded_sam/detection_result.py b/invokeai/backend/grounded_sam/detection_result.py new file mode 100644 index 0000000000..40e4254385 --- /dev/null +++ b/invokeai/backend/grounded_sam/detection_result.py @@ -0,0 +1,41 @@ +from dataclasses import dataclass +from typing import Any, Optional + +import numpy.typing as npt + + +@dataclass +class BoundingBox: + """Bounding box helper class.""" + + xmin: int + ymin: int + xmax: int + ymax: int + + def to_box(self) -> list[int]: + """Convert to the array notation expected by SAM.""" + return [self.xmin, self.ymin, self.xmax, self.ymax] + + +@dataclass +class DetectionResult: + """Detection result from Grounding DINO or Grounded SAM.""" + + score: float + label: str + box: BoundingBox + mask: Optional[npt.NDArray[Any]] = None + + @classmethod + def from_dict(cls, detection_dict: dict[str, Any]): + return cls( + score=detection_dict["score"], + label=detection_dict["label"], + box=BoundingBox( + xmin=detection_dict["box"]["xmin"], + ymin=detection_dict["box"]["ymin"], + xmax=detection_dict["box"]["xmax"], + ymax=detection_dict["box"]["ymax"], + ), + )