Rename labelme.ai -> labelme._automation

labelme/ai/text_to_annotation.py → labelme/_automation/bbox_from_text.py

@@ -2,11 +2,12 @@
 import time
 
 import numpy as np
+import numpy.typing as npt
 import osam
 from loguru import logger
 
 
-def get_rectangles_from_texts(
+def get_bboxes_from_texts(
     model: str, image: np.ndarray, texts: list[str]
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     request: osam.types.GenerateRequest = osam.types.GenerateRequest(
@@ -23,18 +24,18 @@ def get_rectangles_from_texts(
         f"Requesting with model={model!r}, image={(image.shape, image.dtype)}, "
         f"prompt={request.prompt!r}"
     )
-    t_start = time.time()
+    t_start: float = time.time()
     response: osam.types.GenerateResponse = osam.apis.generate(request=request)
 
-    num_annotations = len(response.annotations)
+    num_annotations: int = len(response.annotations)
     logger.debug(
         f"Response: num_annotations={num_annotations}, "
         f"elapsed_time={time.time() - t_start:.3f} [s]"
     )
 
-    boxes: np.ndarray = np.empty((num_annotations, 4), dtype=np.float32)
-    scores: np.ndarray = np.empty((num_annotations,), dtype=np.float32)
-    labels: np.ndarray = np.empty((num_annotations,), dtype=np.int32)
+    boxes: npt.NDArray[np.float32] = np.empty((num_annotations, 4), dtype=np.float32)
+    scores: npt.NDArray[np.float32] = np.empty((num_annotations,), dtype=np.float32)
+    labels: npt.NDArray[np.float32] = np.empty((num_annotations,), dtype=np.int32)
     for i, annotation in enumerate(response.annotations):
         boxes[i] = [
             annotation.bounding_box.xmin,
@@ -48,16 +49,18 @@ def get_rectangles_from_texts(
     return boxes, scores, labels
 
 
-def non_maximum_suppression(
+def nms_bboxes(
     boxes: np.ndarray,
     scores: np.ndarray,
     labels: np.ndarray,
     iou_threshold: float,
     score_threshold: float,
     max_num_detections: int,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    num_classes = np.max(labels) + 1
-    scores_of_all_classes = np.zeros((len(boxes), num_classes), dtype=np.float32)
+    num_classes: int = max(labels) + 1
+    scores_of_all_classes: npt.NDArray[np.float32] = np.zeros(
+        (len(boxes), num_classes), dtype=np.float32
+    )
     for i, (score, label) in enumerate(zip(scores, labels)):
         scores_of_all_classes[i, label] = score
     logger.debug(f"Input: num_boxes={len(boxes)}")
@@ -72,14 +75,14 @@ def non_maximum_suppression(
     return boxes, scores, labels
 
 
-def get_shapes_from_annotations(
+def get_shapes_from_bboxes(
     boxes: np.ndarray, scores: np.ndarray, labels: np.ndarray, texts: list[str]
 ) -> list[dict]:
     shapes: list[dict] = []
     for box, score, label in zip(boxes.tolist(), scores.tolist(), labels.tolist()):
-        text = texts[label]
+        text: str = texts[label]
         xmin, ymin, xmax, ymax = box
-        shape = {
+        shape: dict = {
             "label": text,
             "points": [[xmin, ymin], [xmax, ymax]],
             "group_id": None,

labelme/ai/_utils.py → labelme/_automation/polygon_from_mask.py

@@ -1,24 +1,27 @@
 import imgviz
 import numpy as np
+import numpy.typing as npt
 import skimage
 from loguru import logger
 
 
-def _get_contour_length(contour):
-    contour_start = contour
-    contour_end = np.r_[contour[1:], contour[0:1]]
+def _get_contour_length(contour: npt.NDArray[np.float32]) -> float:
+    contour_start: npt.NDArray[np.float32] = contour
+    contour_end: npt.NDArray[np.float32] = np.r_[contour[1:], contour[0:1]]
     return np.linalg.norm(contour_end - contour_start, axis=1).sum()
 
 
-def compute_polygon_from_mask(mask):
-    contours = skimage.measure.find_contours(np.pad(mask, pad_width=1))
+def compute_polygon_from_mask(mask: npt.NDArray[np.bool_]) -> npt.NDArray[np.float32]:
+    contours: npt.NDArray[np.float32] = skimage.measure.find_contours(
+        np.pad(mask, pad_width=1)
+    )
     if len(contours) == 0:
         logger.warning("No contour found, so returning empty polygon.")
         return np.empty((0, 2), dtype=np.float32)
 
-    contour = max(contours, key=_get_contour_length)
-    POLYGON_APPROX_TOLERANCE = 0.004
-    polygon = skimage.measure.approximate_polygon(
+    contour: npt.NDArray[np.float32] = max(contours, key=_get_contour_length)
+    POLYGON_APPROX_TOLERANCE: float = 0.004
+    polygon: npt.NDArray[np.float32] = skimage.measure.approximate_polygon(
         coords=contour,
         tolerance=np.ptp(contour, axis=0).max() * POLYGON_APPROX_TOLERANCE,
     )

labelme/app.py

@@ -18,7 +18,7 @@
 from PyQt5.QtCore import Qt
 
 from labelme import __appname__
-from labelme import ai
+from labelme._automation import bbox_from_text
 from labelme.config import get_config
 from labelme.label_file import LabelFile
 from labelme.label_file import LabelFileError
@@ -1011,7 +1011,7 @@ def status(self, message, delay=5000):
 
     def _submit_ai_prompt(self, _) -> None:
         texts = self._ai_prompt_widget.get_text_prompt().split(",")
-        boxes, scores, labels = ai.get_rectangles_from_texts(
+        boxes, scores, labels = bbox_from_text.get_bboxes_from_texts(
             model="yoloworld",
             image=utils.img_qt_to_arr(self.image)[:, :, :3],
             texts=texts,
@@ -1033,7 +1033,7 @@ def _submit_ai_prompt(self, _) -> None:
             scores = np.r_[scores, [1.01]]
             labels = np.r_[labels, [texts.index(shape.label)]]
 
-        boxes, scores, labels = ai.non_maximum_suppression(
+        boxes, scores, labels = bbox_from_text.nms_bboxes(
             boxes=boxes,
             scores=scores,
             labels=labels,
@@ -1047,7 +1047,7 @@ def _submit_ai_prompt(self, _) -> None:
         scores = scores[keep]
         labels = labels[keep]
 
-        shape_dicts: list[dict] = ai.get_shapes_from_annotations(
+        shape_dicts: list[dict] = bbox_from_text.get_shapes_from_bboxes(
             boxes=boxes,
             scores=scores,
             labels=labels,

labelme/widgets/canvas.py

@@ -9,8 +9,7 @@
 
 import osam
 import numpy as np
-import labelme.ai
-from labelme.ai._utils import compute_polygon_from_mask
+from labelme._automation import polygon_from_mask
 import labelme.utils
 from labelme.shape import Shape
 
@@ -782,7 +781,9 @@ def paintEvent(self, event: QtGui.QPaintEvent) -> None:
             p.end()
             return
 
-        points = compute_polygon_from_mask(mask=response.annotations[0].mask)
+        points = polygon_from_mask.compute_polygon_from_mask(
+            mask=response.annotations[0].mask
+        )
         if len(points) < 2:
             p.end()
             return