Feature/python rotation (#30)

* add equivalent python image rotation

* attempt to integrate rotation with instance segmentation

Author: nathanhhughes

semantic_inference/python/semantic_inference/image_rotator.py

@@ -0,0 +1,53 @@
+"""Class for performing inference on a rotated image."""
+
+import enum
+
+import numpy as np
+
+
+class RotationType(enum.Enum):
+    """Type of rotation to apply to image."""
+
+    ROTATE_90_CLOCKWISE = "ROTATE_90_CLOCKWISE"
+    ROTATE_180 = "ROTATE_180"
+    ROTATE_90_COUNTERCLOCKWISE = "ROTATE_90_COUNTERCLOCKWISE"
+    NONE = "none"
+
+
+class ImageRotator:
+    """Class to apply and unapply image rotations."""
+
+    def __init__(self, mode):
+        """Construct an image rotator with the transformation to apply for inference."""
+        self._mode = mode
+        match self._mode:
+            case RotationType.ROTATE_90_CLOCKWISE:
+                self._forward_iters = 3
+                self._reverse_iters = 1
+            case RotationType.ROTATE_180:
+                self._forward_iters = 2
+                self._reverse_iters = 2
+            case RotationType.ROTATE_90_COUNTERCLOCKWISE:
+                self._forward_iters = 1
+                self._reverse_iters = 3
+            case _:
+                pass
+
+    @property
+    def mode(self):
+        """Get the rotation mode."""
+        return self._mode
+
+    def rotate(self, img):
+        """Rotate an image to the correct orientation to run inference."""
+        if self._mode == RotationType.NONE:
+            return img
+
+        return np.rot90(img, k=self._forward_iters)
+
+    def derotate(self, img):
+        """Unapply rotation to an image after running inference."""
+        if self._mode == RotationType.NONE:
+            return img
+
+        return np.rot90(img, k=self._reverse_iters)

semantic_inference/python/semantic_inference/models/instance_segmenter.py

@@ -38,6 +38,8 @@
 from spark_config import Config, config_field
 from torch import nn
 
+from semantic_inference.image_rotator import ImageRotator, RotationType
+
 
 def _map_opt(values, f):
     return {k: v if v is None else f(v) for k, v in values.items()}
@@ -52,30 +54,7 @@ class Results:
     boxes: torch.Tensor  # (n, 4) xyxy format, torch.float32
     categories: torch.Tensor  # (n,), torch.float32/int64 (doesn't matter)
     confidences: torch.Tensor  # (n,), torch.float32
-    img_shape: tuple[int, int]
-
-    @property
-    def instance_seg_img(self):
-        """
-        Convert segmentation results to instance segmentation image.
-        Each pixel value encodes both category id and instance id.
-        First 16 bits are category id, last 16 bits are instance id.
-        """
-        if self.masks is None:
-            return np.zeros(self.img_shape)
-
-        masks = self.masks.cpu().numpy()
-        category_ids = self.categories.cpu().numpy()
-        img = np.zeros(masks[0].shape, dtype=np.uint32)
-        for i in range(masks.shape[0]):
-            category_id = int(category_ids[i])  # category id are 0-indexed
-            instance_id = i + 1  # instance ids are 1-indexed
-            combined_id = (
-                category_id << 16
-            ) | instance_id  # combine into single uint32
-            img[masks[i, ...] > 0] = combined_id
-
-        return img
+    instances: np.ndarray  # (H, W, c) np.uint32 (result image)
 
     def cpu(self):
         """Move results to CPU."""
@@ -94,6 +73,7 @@ class InstanceSegmenterConfig(Config):
 
     # relevant configs (model path, model weights) for the model
     instance_model: Any = config_field("instance_model", default="yolo-seg")
+    rotation_type: str = "none"
 
 
 class InstanceSegmenter(nn.Module):
@@ -106,6 +86,7 @@ def __init__(self, config):
         self._canary_param = nn.Parameter(torch.empty(0))
 
         self.config = config
+        self._rotator = ImageRotator(RotationType(config.rotation_type))
         self.segmenter = self.config.instance_model.create()
 
     def eval(self):
@@ -134,7 +115,7 @@ def segment(self, rgb_img, is_rgb_order=True):
             Encoded image
         """
         img = rgb_img if is_rgb_order else rgb_img[:, :, ::-1].copy()
-        return self(img)
+        return self._rotator.derotate(img)
 
     @property
     def device(self):
@@ -156,12 +137,28 @@ def forward(self, rgb_img):
         Returns:
             Encoded image
         """
-        categories, masks, boxes, confidences, img_shape = self.segmenter(rgb_img)
+        rotated = self._rotator.rotate(rgb_img)
+        categories, masks, boxes, confidences = self.segmenter(rotated)
+
+        if self.masks is None:
+            instances = np.zeros(rgb_img.shape[:2])
+        else:
+            instances = np.zeros(masks[0].shape, dtype=np.uint32)
+            masks = self.masks.cpu().numpy()
+            category_ids = self.categories.cpu().numpy()
+            for i in range(masks.shape[0]):
+                category_id = int(category_ids[i])  # category id are 0-indexed
+                instance_id = i + 1  # instance ids are 1-indexed
+                # combine into single uint32
+                combined_id = (category_id << 16) | instance_id
+                instances[masks[i, ...] > 0] = combined_id
+
+            instances = self._rotator.derotate(instances)
 
         return Results(
             masks=masks,
             boxes=boxes,
             categories=categories,
             confidences=confidences,
-            img_shape=img_shape,
+            instances=instances,
         )

semantic_inference/python/semantic_inference/models/wrappers.py

@@ -381,7 +381,7 @@ def forward(self, img):
             iou=self.config.overlap_merge_iou,
         )[0]  # assume batch size 1
         if result.masks is None:
-            return None, None, None, None, (img.shape[0], img.shape[1])
+            return None, None, None, None
 
         categories = result.boxes.cls.to(torch.int).cpu()  # int8
         masks = result.masks.data.to(torch.bool).cpu()
@@ -397,7 +397,7 @@ def forward(self, img):
             boxes = boxes[size_indices]
             confidences = confidences[size_indices]
 
-        return categories, masks, boxes, confidences, (masks.shape[1], masks.shape[2])
+        return categories, masks, boxes, confidences
 
 
 @dataclasses.dataclass
@@ -569,7 +569,7 @@ def forward(self, img):
 
         # if nothing detected
         if boxes.shape[0] == 0:
-            return None, None, None, None, (img.shape[0], img.shape[1])
+            return None, None, None, None
 
         # process the box prompt for SAM 2
         h, w, _ = img.shape
@@ -633,7 +633,7 @@ def forward(self, img):
         # use xyxy boxes
         boxes = torch.tensor(input_boxes)
 
-        return categories, masks, boxes, confidences, (masks.shape[1], masks.shape[2])
+        return categories, masks, boxes, confidences
 
 
 @register_config(

semantic_inference/python/test/test_image_rotator.py

@@ -0,0 +1,66 @@
+"""Unit tests for image rotator."""
+
+import numpy as np
+
+from semantic_inference.image_rotator import ImageRotator, RotationType
+
+
+def test_no_rotation():
+    """Test that no rotation produces the same matrix."""
+    A = np.arange(24).reshape((3, 4, 2))
+    rotator = ImageRotator(RotationType.NONE)
+
+    assert (rotator.rotate(A) == A).all()
+    assert (rotator.derotate(A) == A).all()
+
+
+def test_90cw_rotation():
+    """Test that no rotation produces the same matrix."""
+    A = np.arange(24).reshape((3, 4, 2))
+
+    rotator = ImageRotator(RotationType.ROTATE_90_CLOCKWISE)
+    rotated = rotator.rotate(A)
+    expected = np.array(
+        [
+            [[16, 17], [8, 9], [0, 1]],
+            [[18, 19], [10, 11], [2, 3]],
+            [[20, 21], [12, 13], [4, 5]],
+            [[22, 23], [14, 15], [6, 7]],
+        ]
+    )
+
+    assert rotated.shape == (4, 3, 2)
+    assert (rotated == expected).all()
+    assert (rotator.derotate(rotated) == A).all()
+
+
+def test_180_rotation():
+    """Test that no rotation produces the same matrix."""
+    A = np.arange(24).reshape((3, 4, 2))
+
+    rotator = ImageRotator(RotationType.ROTATE_180)
+    rotated = rotator.rotate(A)
+
+    assert rotated.shape == A.shape
+    assert (rotated != A).any()
+    assert (rotator.derotate(rotated) == A).all()
+
+
+def test_90ccw_rotation():
+    """Test that no rotation produces the same matrix."""
+    A = np.arange(24).reshape((3, 4, 2))
+
+    rotator = ImageRotator(RotationType.ROTATE_90_COUNTERCLOCKWISE)
+    rotated = rotator.rotate(A)
+    expected = np.array(
+        [
+            [[6, 7], [14, 15], [22, 23]],
+            [[4, 5], [12, 13], [20, 21]],
+            [[2, 3], [10, 11], [18, 19]],
+            [[0, 1], [8, 9], [16, 17]],
+        ]
+    )
+
+    assert rotated.shape == (4, 3, 2)
+    assert (expected == rotated).all()
+    assert (rotator.derotate(rotated) == A).all()

semantic_inference_ros/app/instance_segmentation_node

@@ -85,10 +85,8 @@ class InstanceSegmentationNode(Node):
             self.get_logger().debug("get image size: " + str(ret.img_shape))
             self.get_logger().debug(f"result category type: {ret.categories.dtype}")
 
-        # get 32 bit instance + cateogry img, all 0 if no box detected
-        instance_seg_img = ret.instance_seg_img
         # Convert to int32 to match 32SC1 encoding expected by cv_bridge
-        instance_seg_img = instance_seg_img.astype(np.int32)
+        instance_seg_img = ret.instances.astype(np.int32)
         msg = Conversions.to_image_msg(header, instance_seg_img, encoding="32SC1")
         self._pub.publish(msg)
         self._last_inference_time = self.get_clock().now().nanoseconds * 1e-9