9.1. Augmented VR Module

import inspect
from viper_toolkit import Dissect
from scripts import augmented_vr
from scripts.augmented_vr import AugmentedVRModule
source = inspect.getsource(AugmentedVRModule)
print (source)
class AugmentedVRModule(object):
    
    def __init__(self):
        self.name = NameManager()
        #self.dynamic_updates()
        self.pose_estimations = InferenceResults()
        self.setup_parameters()
        self.setup_ros()
        self.setup_scene_segmentation()
        self.loop()

    def setup_parameters(self):
        # Adds all of our parameters to an updatable list. If ROS cannot
        # find the parameter, it will use the default provided.
        self.parameters = ParameterManager()
        
        self.parameters.add(Parameter("width", "/inland_ir_cam/width", default = 800, dynamic = True))
        self.parameters.add(Parameter("height", "/inland_ir_cam/height", 600, True))
        self.parameters.add(Parameter("updates", f"{self.name.name}/dynamic", True, True))
        self.parameters.add(Parameter("image_output", f"{self.name.name}/image_output", True, True))
        self.parameters.add(Parameter("alpha", rospy.search_param('alpha'), 0.3, True))
        self.parameters.add(Parameter("threshold", rospy.search_param('threshold'), 0.1, True))
        self.parameters.add(Parameter("segmentationmodel", f"{self.name.name}/segmentation", True, True))
        self.parameters.add(Parameter("posedetectionmodel", f"{self.name.name}/posedetection", True, True))
        self.parameters.add(Parameter("rate", f"{self.name.name}/rate", 25, True))
        
    def setup_ros(self):
        rospy.init_node('Augmented_VR', log_level = rospy.DEBUG)
        self.bridge = CvBridge()
        self.setup_camera()
        self.setup_pose_detector()
        self.setup_scene_segmentation()
        self.setup_vr_publisher()
        
    def setup_vr_publisher(self):
        self.pub = rospy.Publisher(
            'Augmented_VR/stream',
            Image,
            queue_size=1
            )
            
    def setup_camera(self):
        rospy.Subscriber("image_server/image",
            Image, 
            self.image_callback, 
            queue_size=1
            )
        rospy.loginfo("Camera Online")

        self.image = Image()
        self.new_image = True

        self.canvas = np.zeros((
                    self.parameters.height,
                    self.parameters.width, 
                    3
                    ), dtype = "uint8"
                    )

    def setup_pose_detector(self):
        rospy.Subscriber(
            'model_output/pose_articulator/human_graph_nodes',
            InferenceResults,
            self.pose_estimation_callback,
            queue_size=1
            )
        self.pose_estimations = InferenceResults()
        self.new_graph = True
        self.graph = []
        rospy.loginfo("Pose detector subscription active")

    def setup_scene_segmentation(self):
        self.mask = InferenceResults()
        self.new_mask = True
        rospy.Subscriber(
            'model_output/scene_segmentation/segmentation_mask',
            InferenceResults,
            self.scene_segmentation_callback,
            queue_size=1
            )
        self.mask = InferenceResults()
        self.colormap = np.array([[68, 1, 84], [48, 103, 141], [53, 183, 120], [199, 216, 52]])
        rospy.loginfo("Scene Segmentation subscription active")

    def pose_estimation_callback(self, msg):
        self.pose_estimations = msg
        self.new_graph = True
        print (self.pose_estimations)
        rospy.loginfo("[POS] New Coordinates Received")

    def scene_segmentation_callback(self, msg):
        self.mask = msg
        self.new_mask = True
        print (self.mask)
        rospy.loginfo("[SEG] New Masks Received")
   
    def image_callback(self, msg):
        self.image = msg
        self.new_image = True
        rospy.loginfo("[IMG] New Image Received")

    def decode_inference_results(self, results, model: str = None):
        if model == None:
            name = self.name.abv
        else:
            name = model
        self.timer.lap("Decode Inference")
        rospy.logdebug(f"Decoding {name}")
        try:
            structure = results.structure
            print(name, structure)
        except:
            rospy.logerr(
            f"[{name}] Cannot convert {type(self.results.structure)} to structure"
            )
            return
        print ("------------------------------------------")
        try:
            inferences = np.asarray(results.inferences)
            print(inferences)
            rospy.loginfo(f"[{name}] Transformed shape: {inferences.shape}")
        except:
            rospy.logerr(f'[{name}] Cannot convert to numpy array')
            return
            
        try:
            inference_array = inferences.reshape(results.structure)
            print(inference_array)
        except:
            rospy.logerr(f"[{name}] Cannot wrangle data blob")
            return
            
        #self.timer.time("Decode Inference", name=name)
        
        return inference_array
            
    def update_segmentation_graph(self):
    # This function tasks a message file (self.mask) and decodes it into
    # the original array, outputing the "segmentation graph". Note that
    # this output can be used on its own as an image.
        print (self.mask)
        #self.timer.lap("Segmentation Process")
        #try:
        inference_array = self.decode_inference_results(
                self.mask, 
                model="SEG"
                )
        #rospy.loginfo(f'inference_array is: {type(inference_array)} {inference_array.shape}')
        #except:
        #    rospy.logerr(f"[SEG] Cannot decode results")
        #    return
        
        #try:
            #self.timer.lap("Segmentation Model")
            #rospy.logdebug(f"[{self.name.abv}] Conversion Shape: %s",inference_array.shape)
        print(type(inference_array))
        mask = segmentation_map_to_image(inference_array, self.colormap)
            #rospy.logdebug(f"[{self.name.abv}] Resizing Mask")
        self.segmentation_graph = cv2.resize(inference_array, (self.image.width, self.image.height))
        self.logdebug(f'seg graph shape {self.segmentation_graph.shape}')
            #rospy.logdebug(f"[{self.name.abv}] Masking Image")
            #self.timer.time("Segmentation Model", name="SEG")
            #self.timer.time("Segmentation Process", name="SEG")
        return
        #except:
        #    rospy.logerr(f"[SEG] Cannot resize mask at point A")
            #self.timer.time("Segmentation Model", name="SEG")
            #self.timer.time("Segmentation Process", name="SEG")
        #    return

    def combine_segmentation_image(self):
        # This function combines the current segmentation_graph 
        # and the current converted image, and saves it as the 
        # current output (canvas)
        alpha = 0.3
        self.timer.lap("Image Overlay")
        #try:
        self.canvas = cv2.addWeighted(
            self.segmentation_graph, 
            alpha, #self.parameters.alpha, 
            self.img, 
            1 - alpha, 
            0)
        #except:
        #    rospy.logerr(f"[SEG] Cannot combine mask and image at B")
            
        # Regardless of success we will update our timer.
        #self.timer.time("Image Overlay", name = "SEG")
        return 

    def calculate_pose_graph(self):
        # This process takes the current pose estimation msg
        # (self.pose_estimates) and converts them to the 
        # pose graph nodes and edges. It then saves these
        # to the current nodes and edges (self.graph)

        # Convert the message arrays
        #self.timer.lap("Graph Nodes Edges")
        rospy.logdebug("Decoding Pose Graph Results")
        #try:
        self.graph = self.decode_inference_results(
            self.pose_estimations, 
            model="POS"
            )
        #except:
            # If this fails there is nothing to return.
        #    rospy.logerr(f"[POS] Cannot decode results")
            
        #self.timer.time("Graph Nodes Edges")

    def draw_graph(self):
        # Draw the Poses on the provided image 'image' and return
        # the composite 'graph'
        self.timer.lap("Draw Poses")
        try:
            self.canvas = draw_poses(
                img = self.canvas, 
                poses = self.graph, 
                point_score_threshold = 0.1) #elf.parameters.threshold

        except:
            rospy.logerr(f"[POS] Cannot draw Poses")
        self.timer.time("Draw Poses", name="POS")

    def action_loop(self):

        # First, we need to decide if we are using an existing 
        # segmentation mask, or creating a new segmentation. If there
        # are new masks, we will update the self.segmentation_graph
        # and then remove our flag since we can use this until the new
        # results are delivered.
        if self.new_mask:
            self.update_segmentation_graph()
            self.new_mask = False
        
        # Next we will decide if we are recalculating the existing
        # graph nodes at self.graph.
        if self.new_graph:
            self.calculate_pose_graph()
            self.new_graph = False
        
        # If we are combining the segmentation and the images, then we
        # first need to mask the new images, but we only need to do
        # this if we need to combine the images and the inferences.
        temp = True
        if temp: #self.parameters.image_output:
            if self.new_image:
                try:
                    # Convert the camera image to ROS
                    self.img = self.bridge.imgmsg_to_cv2(
                        self.image, 
                        desired_encoding="bgr8"
                        )
                    self.new_graph = False
                except CvBridgeError as e:
                    # if this fails then give us the last good image.
                    rospy.logerr("[IMG] Error converting ROS msg to image.")
                    print(e)
                    
            if temp: #self.parameters.segmentationmodel:
                try:
                    # We can now combine them on the 'self.canvas' layer.
                    self.combine_segmentation_image()
                except:
                    rospy.logerr("[SEG] Error proccessing scene")
        else:
            if self.parameters.segmentation_model:
            # If we are not using the camera image output then we
            # will just replace the self.image with the segmentation
            # graph image.
            
                self.canvas = self.segmentation_graph
            else:
                self.canvas = np.zeros((
                    self.height, 
                    self.width, 
                    3
                    ), dtype = "uint8"
                    )
            
        # Next we will draw our the current pose graph nodes and edges
        # onto our canvas
        
        #if self.parameters.posedetectionmodel == True:
        try:
            self.draw_graph()
        except:
            rospy.logerr("[POS] Error drawing pose detection graph")

        # Finally we will convert our image canvase back to the 
        # ROS message format.
        try:
            image_ros = self.bridge.cv2_to_imgmsg(self.canvas, 'bgr8')
            return image_ros
                
        except (CvBridgeError, TypeError, UnboundLocalError) as e:
            # If this conversion failes we will simply provide the last
            # good ROS image message since sometimes packages are lost
            rospy.logerr("[IMG]Error converting to ROS msg")
            print (e)
            return self.image

    def loop(self):

        rate = rospy.Rate(self.parameters.rate)
        while not rospy.is_shutdown():  
            # We have created a parameter for dynamically updateding an
            # active node. These will update if the global value changes
            #if self.parameters.updates == True:
            #    self.parameters.update()
            self.timer = ProcessTimer(abv = self.name.abv)
            ros_image = self.action_loop()
            self.pub.publish(ros_image)
            rate.sleep()