initial_tune

segment_anything_ui/utils/shapes.py

@@ -0,0 +1,53 @@
+import dataclasses
+
+import cv2
+import numpy as np
+from PySide6.QtCore import QPoint
+from PySide6.QtGui import QPolygon
+
+
+@dataclasses.dataclass
+class BoundingBox:
+    xstart: float | int
+    ystart: float | int
+    xend: float | int = -1.
+    yend: float | int = -1.
+
+    def to_numpy(self):
+        return np.array([self.xstart, self.ystart, self.xend, self.yend])
+
+    def scale(self, sx, sy):
+        return BoundingBox(
+            xstart=self.xstart * sx,
+            ystart=self.ystart * sy,
+            xend=self.xend * sx,
+            yend=self.yend * sy
+        )
+
+    def to_int(self):
+        return BoundingBox(
+            xstart=int(self.xstart),
+            ystart=int(self.ystart),
+            xend=int(self.xend),
+            yend=int(self.yend)
+        )
+
+@dataclasses.dataclass
+class Polygon:
+    points: list = dataclasses.field(default_factory=list)
+
+    def to_numpy(self):
+        return np.array(self.points).reshape(-1, 2)
+
+    def to_mask(self, num_rows, num_cols):
+        mask = np.zeros((num_rows, num_cols))
+        mask = cv2.fillPoly(mask, pts=[self.to_numpy(), ], color=255)
+        return mask
+
+    def is_plotable(self):
+        return len(self.points) > 3
+
+    def to_qpolygon(self):
+        return QPolygon([
+            QPoint(x, y) for x, y in self.points
+        ])