Method update

TEACh/task_base/object_tracker.py:213–491

rgb: RGB image depth: depth image camX0_T_camX: rotation pose matrix to go from reference frame to current frame

(
        self, 
        rgb, 
        depth, 
        camX0_T_camX, 
        return_det_dict=False, 
        use_gt=False, 
        target_object=None, 
        only_keep_target=False,
        vis=None, 
        return_features=False,
        force_update_odin=False,
        )

                                )).double()

    def update(
        self, 
        rgb, 
        depth, 
        camX0_T_camX, 
        return_det_dict=False, 
        use_gt=False, 
        target_object=None, 
        only_keep_target=False,
        vis=None, 
        return_features=False,
        force_update_odin=False,
        ):
        '''
        rgb: RGB image
        depth: depth image
        camX0_T_camX: rotation pose matrix to go from reference frame to current frame
        '''

        if args.simulate_actions:
            # do not update if simulating
            return

        out = {}

        if not (depth.shape[-2]==self.W and depth.shape[-1]==self.H):
            depth = cv2.resize(depth, (self.W, self.H), interpolation=cv2.INTER_NEAREST)
        

        ####### DETECT OBJECTS ##########
        if self.use_gt_objecttrack:
            pred_scores, pred_labels, pred_boxes_or_masks, centroids_gt, attributes_gt = self.get_objects_gt(self.controller, depth)
            if not args.use_gt_centroids:
                if len(pred_boxes_or_masks)>0:
                    depth_ = torch.from_numpy(depth).cuda().unsqueeze(0).unsqueeze(0)
                    xyz = utils.geom.depth2pointcloud(depth_, torch.from_numpy(self.pix_T_camX).cuda().unsqueeze(0).float())
                    xyz_origin = utils.geom.apply_4x4(camX0_T_camX.cuda().float(), xyz).squeeze().cpu().numpy()
                    xyz_origin = xyz_origin.reshape(1,self.W,self.H,3) 
                else:
                    centroids = []
        elif self.use_odin:            
            if not force_update_odin:
                depth = self.navigation.task.get_observations()["depth"].copy()
                depth[depth < 0.5] = 0.
                self.odin_input_dict['images'].append(rgb.copy())
                self.odin_input_dict['depths'].append(depth)
                origin_T_camX = self.world_t_weird @ camX0_T_camX
                self.odin_input_dict['poses'].append(origin_T_camX)
                self.odin_input_dict['intrinsics'].append(self.pix_T_camX)
                self.steps_since_odin_update += 1
            if self.steps_since_odin_update>args.odin_update_frequency or force_update_odin:
                pred_boxes_or_masks, pred_labels, pred_scores, centroids_odin, crops_odin = self.multiview_detector.get_masks(
                    self.odin_input_dict, 
                    target_class=target_object,
                    id_to_mapped_id=self.id_to_mapped_id,
                    )
                self.steps_since_odin_update = 0
            else: