From 79b9546336eb86f488733f5d0d5d08b3afec72f4 Mon Sep 17 00:00:00 2001 From: Nathan Hughes Date: Thu, 16 Apr 2026 02:16:33 +0000 Subject: [PATCH 1/2] add equivalent python image rotation --- .../semantic_inference/image_rotator.py | 53 +++++++++++++++ .../python/test/test_image_rotator.py | 66 +++++++++++++++++++ 2 files changed, 119 insertions(+) create mode 100644 semantic_inference/python/semantic_inference/image_rotator.py create mode 100644 semantic_inference/python/test/test_image_rotator.py diff --git a/semantic_inference/python/semantic_inference/image_rotator.py b/semantic_inference/python/semantic_inference/image_rotator.py new file mode 100644 index 0000000..19c8243 --- /dev/null +++ b/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) diff --git a/semantic_inference/python/test/test_image_rotator.py b/semantic_inference/python/test/test_image_rotator.py new file mode 100644 index 0000000..546fea8 --- /dev/null +++ b/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() From 547fb279e498b1c6aa9a64fbd7fd8e545e8c9a2b Mon Sep 17 00:00:00 2001 From: Nathan Hughes Date: Thu, 16 Apr 2026 03:05:31 +0000 Subject: [PATCH 2/2] attempt to integrate rotation with instance segmentation --- .../models/instance_segmenter.py | 51 +++++++++---------- .../semantic_inference/models/wrappers.py | 8 +-- .../app/instance_segmentation_node | 4 +- 3 files changed, 29 insertions(+), 34 deletions(-) diff --git a/semantic_inference/python/semantic_inference/models/instance_segmenter.py b/semantic_inference/python/semantic_inference/models/instance_segmenter.py index 7b8137e..44354b0 100644 --- a/semantic_inference/python/semantic_inference/models/instance_segmenter.py +++ b/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, ) diff --git a/semantic_inference/python/semantic_inference/models/wrappers.py b/semantic_inference/python/semantic_inference/models/wrappers.py index 55e7be4..a528b56 100644 --- a/semantic_inference/python/semantic_inference/models/wrappers.py +++ b/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( diff --git a/semantic_inference_ros/app/instance_segmentation_node b/semantic_inference_ros/app/instance_segmentation_node index 372d5db..9e6776e 100755 --- a/semantic_inference_ros/app/instance_segmentation_node +++ b/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