projection.py

from __future__ import annotations

import numpy as np

from numpy import ndarray
from .projection_geometry import PyProjectionGeometry

from rq_interfaces.msg import Projection, ProjectionGeometry
from sensor_msgs.msg import Image
import ros2_numpy


class PyProjection(PyProjectionGeometry):
    def __init__(self, focal_spot_mm: ndarray, detector_postion_mm: ndarray, detector_orientation_quad: ndarray, image: np.ndarray,
                 detector_heigth_mm: float, detector_width_mm: float, frame_id: str = 'object', 
                 focal_spot_orientation_quad: np.ndarray = np.array([0., 0., 0, 1.]),
                 voltage_kv: float = 100., current_ua: float = 100., exposure_time_ms: float = 1000.) -> None:
        super().__init__(focal_spot_mm, detector_postion_mm, detector_orientation_quad, frame_id, focal_spot_orientation_quad)
        self.image = image.astype(np.uint16)
        self.detector_heigth_mm = detector_heigth_mm
        self.detector_width_mm = detector_width_mm
        self.voltage_kv = voltage_kv
        self.current_ua = current_ua
        self.exposure_time_ms = exposure_time_ms

    @classmethod
    def dummy(cls):
        return cls(np.array([0., 100., 0]), 
                   np.array([0., -100., 0]),
                   np.array([1., 0., 0, 1.]),
                   np.zeros((10, 10)),
                   10., 10.)

    @classmethod
    def from_message(cls, msg: Projection):
        focal_spot_mm = np.array([msg.projection_geometry.focal_spot_postion_mm.x,
                                  msg.projection_geometry.focal_spot_postion_mm.y,
                                  msg.projection_geometry.focal_spot_postion_mm.z,])

        detector_center_mm = np.array([msg.projection_geometry.detector_postion_mm.x,
                                       msg.projection_geometry.detector_postion_mm.y,
                                       msg.projection_geometry.detector_postion_mm.z,])
        
        detector_orientation_quad = np.array([msg.projection_geometry.detector_orientation_quad.x,
                                              msg.projection_geometry.detector_orientation_quad.y,
                                              msg.projection_geometry.detector_orientation_quad.z,
                                              msg.projection_geometry.detector_orientation_quad.w])
        
        focal_spot_orientation = np.array([msg.projection_geometry.focal_spot_orientation_quad.x,
                                           msg.projection_geometry.focal_spot_orientation_quad.y,
                                           msg.projection_geometry.focal_spot_orientation_quad.z,
                                           msg.projection_geometry.focal_spot_orientation_quad.w])
        
        
        detector_heigth_mm = msg.detector_heigth_mm
        detector_width_mm = msg.detector_width_mm
        frame_id = msg.projection_geometry.header.frame_id

        image = ros2_numpy.numpify(msg.image)

        voltage_kv = msg.voltage_kv
        current_ua = msg.current_ua
        exposure_time_ms = msg.exposure_time_ms

        return cls(focal_spot_mm, detector_center_mm, detector_orientation_quad, image, 
                   detector_heigth_mm, detector_width_mm, frame_id, focal_spot_orientation,
                   voltage_kv, current_ua, exposure_time_ms)
    
    def as_message(self) -> ProjectionGeometry:
        message = Projection()
        projection_geometry = ProjectionGeometry()

        projection_geometry.focal_spot_postion_mm.x = self.focal_spot_mm[0]
        projection_geometry.focal_spot_postion_mm.y = self.focal_spot_mm[1]
        projection_geometry.focal_spot_postion_mm.z = self.focal_spot_mm[2]

        projection_geometry.detector_postion_mm.x = self.detector_postion_mm[0]
        projection_geometry.detector_postion_mm.y = self.detector_postion_mm[1]
        projection_geometry.detector_postion_mm.z = self.detector_postion_mm[2]

        projection_geometry.detector_orientation_quad.x = self.detector_orientation_quad[0]
        projection_geometry.detector_orientation_quad.y = self.detector_orientation_quad[1]
        projection_geometry.detector_orientation_quad.z = self.detector_orientation_quad[2]
        projection_geometry.detector_orientation_quad.w = self.detector_orientation_quad[3]

        projection_geometry.focal_spot_orientation_quad.x = self.focal_spot_orientation_quad[0]
        projection_geometry.focal_spot_orientation_quad.y = self.focal_spot_orientation_quad[1]
        projection_geometry.focal_spot_orientation_quad.z = self.focal_spot_orientation_quad[2]
        projection_geometry.focal_spot_orientation_quad.w = self.focal_spot_orientation_quad[3]

        message.projection_geometry = projection_geometry
        message.image = ros2_numpy.msgify(Image, self.image, 'mono16')

        message.detector_heigth_mm = self.detector_heigth_mm
        message.detector_width_mm = self.detector_width_mm

        message.projection_geometry.header.frame_id = self.frame_id

        message.voltage_kv = self.voltage_kv
        message.current_ua = self.current_ua
        message.exposure_time_ms = self.exposure_time_ms

        return message
    
    @property
    def detector_heigth_px(self) -> int:
        return self.image.shape[0]
    
    @property
    def detector_width_px(self) -> int:
        return self.image.shape[1]
    
    @property
    def pixel_pitch_x_mm(self) -> float:
        return self.detector_width_mm / self.detector_width_px
    
    @property
    def pixel_pitch_y_mm(self) -> float:
        return self.detector_heigth_mm / self.detector_heigth_px