region_of_intrest.py

import numpy as np

from rq_interfaces.msg import RegionOfIntrest
from visualization_msgs.msg import Marker


class PyRegionOfIntrest():
    def __init__(self, center_points_mm: np.ndarray, dimensions_mm: np.ndarray, frame_id: str = 'object', 
                 resolution_mm: np.ndarray = np.array([0.1, 0.1, 0.1])):
        self.center_points_mm = center_points_mm
        self.dimensions_mm = dimensions_mm
        self.frame_id = frame_id
        self.resolution_mm = resolution_mm

    @classmethod
    def dummy(cls):
        return cls((np.random.random((3, )) - 0.5) * 20., 
                   (np.random.random((3, )) - 0.5) * 10.)

    @classmethod
    def from_message(cls, msg: RegionOfIntrest):
        center_points_mm = list()
        dimensions_mm = list()

        for roi in msg.region_of_intrest_stack.markers:
            roi: Marker
            center_points_mm.append(
                np.array([roi.pose.position.x,
                          roi.pose.position.y,
                          roi.pose.position.z]))
            
            dimensions_mm.append(
                np.array([roi.scale.x,
                          roi.scale.y,
                          roi.scale.z]))
            
            frame = roi.header.frame_id
        resolution_mm = np.array([msg.resolution.x,
                                  msg.resolution.y,
                                  msg.resolution.z])
        
        return cls(np.array(center_points_mm), np.array(dimensions_mm), frame, resolution_mm)
    
    @property
    def number_of_rois(self) -> int:
        return self.center_points_mm.shape[0]
    
    def as_message(self) -> RegionOfIntrest:
        message = RegionOfIntrest()
        roi_list = list()

        for i in range(self.number_of_rois):
            roi = Marker()

            roi.pose.position.x = self.center_points_mm[i][0]
            roi.pose.position.y = self.center_points_mm[i][1]
            roi.pose.position.z = self.center_points_mm[i][2]

            roi.scale.x = self.dimensions_mm[i][0]
            roi.scale.y = self.dimensions_mm[i][1]
            roi.scale.z = self.dimensions_mm[i][2]

            roi.header.frame_id = self.frame_id

            roi_list.append(roi)            
        
        message.region_of_intrest_stack.markers = roi_list

        message.resolution.x = self.resolution_mm[0]
        message.resolution.x = self.resolution_mm[1]
        message.resolution.x = self.resolution_mm[2]

        return message