MBAPose: Mask and Bounding-Box Aware Pose Estimation of Surgical Instruments with Photorealistic Domain Randomization

The biggest challenge in the use of synthetic images for training deep-learning algorithms is the so-called domain gap. In the context of this work, domain gap is considered as the difference in the performance of the algorithm trained with only synthetic images when it is evaluated against real images. Bridging the domain gap is an open research problem and is being investigated on many fronts, such as in improved rendering, domain adaptation, and domain randomization strategies.

In this work, we build upon our previous work [6] and improve from the pose estimation of the instrument’s shaft to pose estimation of the instruments’ tips. Moreover, we (1) improved the PBR rendering strategy, (2) studied the effects of different DR strategies and proposed two new randomized components, (3) used an improved pose estimation model, and (4) discussed how to practically obtain a real instrument pose dataset and provided results of validation using real data.

To the best of our knowledge, this is the first work to combine those strategies and validate them in the pose estimation of rigid surgical instruments.

The proposed realistic rendering strategy uses elements of PBR and the cut-and-paste method. It calculates light reflection, diffusion, and absorption faithfully with PBR using ray tracing [26] and the concept of microfacet [27] n the same manner as in previous PBR methods [9, 8, 10, 11]. The background is not rendered using PRB; instead, as described in our previous works [6], we render a cube whose faces are real background images. The reflection of the surroundings in the instruments are rendered by tracing photons between the cubic faces and the instruments’ surface.

In a direct comparison with [6], we improved several rendering methods. First, we applied a bidirectional scattering diffusion function (BSDF) with Oren–Nayar reflectance model [28] on the cubic surface to reflect it in various directions. This was necessary because we simplified the surroundings as a cube. A BSDF with GGX²²2GGX itself the name of the distribution and it does not seem to be any particular acronym [29]. microfacet distribution [29] was used on the instruments’ surface to obtain realistic reflection. We positioned a black plane with a hole in front of the camera to reproduce the circular endoscopic view. The sheet was small and placed close to the camera so as to not interfere with the ray tracing.

The examples of synthetic images created with the proposed rendering method are shown is Fig. 3.

DR methods have usually been applied together with classic or simple rendering. We herein combine PBR with DR and propose two further DR strategies.

We use the following previously propose DR strategies: we randomize the position, intensity, number, and type of lights. In addition, the real background images used for the cubic texture are randomly selected. For the background images, we separately evaluated only the real images of the experimental environment, images of the experimental environment and their augmented images, or these plus non-contextual images. As for the surgical instruments, we subdivided the previously proposed texture randomization [22] into color randomization and glossiness randomization for the metallic surface. We were also interested in analyzing the effect of randomizing the textures for each individual instrument separately.

In addition to those well-known DR strategies, we propose a further surface randomization method. It randomizes the appearance instrument’s surface using normal mapping [30]. Normal mapping is a strategy used to change surface appearance of objects by changing the normal vectors of the surface at the rendering time. We randomly created concave and convex surface imperfections by changing the normal maps.

The performance of each strategy in the context of this work is evaluated in the experiments in comprehensive ablation studies.

We propose a new pose estimation model named MBAPose to predict the pose of instruments is aware of their mask and bounding box shape. This model is a combination of two parts, a prediction head and a feature extraction model.

The prediction head, which we call the Box-Aware Head, is one we proposed in [6]. It uses the features obtained from the feature extraction model to simultaneously predict the bounding box, class, and pose. The architecture that internally used self-bounding box prediction and self-class prediction to predict their pose enhanced pose estimation. .

This feature extraction is novel and composed of a double encoder model with mask guidance. The double encoder model itself comprises two smaller encoder models and one decoder model. We propose this feature extractor because our previous model [6] had difficulties in detecting small objects, as shown in [31]. This change was inspired by DSSD [32], that improved the detection of small objects by adding a decoder network. Different from DSSD, we further add a second encoder. The aim is to create a mask-aware feature by adding a loss function with an additional convolution at the head of the decoder model in the training stage. By adding segmentation loss, the model estimates the pose of the objects aware of their contour. With the proposed double encoder, the contour information can be more explicitly used compared with others [33].

One challenge we faced in our prior work [6] was the difficulty in creating a reliable dataset for testing the model with the image and pose information of the real instruments. In this work, we overcame those difficulties and created a test dataset as follows.

The pose of the real instruments was acquired using an optical tracker (Polaris Vega, NDI, Canada) attached to its optical marker at the base of the instruments and endoscope, as shown in Fig. 6 using 3D-printed attachments. However, even small mounting and fabrication errors induce a considerably large error at the tip of the instrument. To calibrate for the mounting error, we propose a practical markerless hand-eye calibration method using the contour of instruments in the view. The method is as follows.

The present hand-eye calibration problem can be represented using homogenous-transformation matrices as $\boldsymbol{AX}=\boldsymbol{ZCB}$, where $\boldsymbol{A}$, $\boldsymbol{B}$, $\boldsymbol{C}$, $\boldsymbol{X}$, and $\boldsymbol{Z}$ denote the transformation from the instrument’s tip pose to the endoscopic lens pose, the instrument’s maker pose with respect to the world, the endoscope’s marker pose with respect to the world, the transformation from the instruments’ maker pose to the instrument’s tip pose, and the transformation from the endoscope’s maker pose to the endoscopic lens pose. $\boldsymbol{B}$ and $\boldsymbol{C}$ are acquired from the optical tracker. We have the ideal values for $\boldsymbol{X}$ and $\boldsymbol{Z}$ using our CAD models; however, they are corrupted by the mounting errors and consequently our calibration target. $\boldsymbol{A}$ is usually estimated with a checkerboard; however, in our case, we do not obtain $\boldsymbol{A}$ explicitly because installing a checkerboard pattern on the tip of the instrument is not feasible. Instead, we indirectly obtain $\boldsymbol{A}$ using the contour of instruments on the images.

We search $\boldsymbol{X}$ and $\boldsymbol{Z}$ by iteratively comparing manually annotated contours on the real images and contours of projected 3D CAD models based on our current estimated of $\boldsymbol{X}$ and $\boldsymbol{Z}$. To do so, we use the following objective function

$$O=\sum_{n}^{N}\left\{\alpha O_{iou}\left(x_{n}^{gt},x_{n}^{proj}\right)+\left(1-\alpha\right)O_{edge}\left(x_{n}^{gt},x_{n}^{proj}\right)\right\},$$

(1)

where $x_{n}^{gt}$ and $x_{n}^{proj}$ are the $n-$th contours on the ground truth and projected CAD model. In detail, $O_{iou}$ is intersection over union (IoU). In addition, $O_{edge}$ is a metric that quantifies edge overlap as

$$O_{edge}\left(x_{gt},x_{proj}\right)=\frac{\sum_{u}\sum_{v}p_{gt}\left(u,v\right)p_{proj}\left(u,v\right)}{\sum_{u}\sum_{v}p_{gt}\left(u,v\right)^{2}},$$

where $p(u,v)$ are the distance transformed pixel values defined

$$p\left(u,v\right)=\left(d_{max}-max\left(d_{nearest}\left(u,v\right),d_{max}\right)\right)^{2},$$

where $d_{nearest}\left(u,v\right)$ is the distance from the nearest edge, and $d_{max}$ is the maximum distance. With this heuristic, we search for, $\boldsymbol{X}$ and $\boldsymbol{Z}$, to maximize $O$ using a ternary search, as shown in Fig. 7.

In this section, we briefly test the performance of our practical hand-eye calibration.

First, we prepared three sequences of endoscopic images paired with the optical tracker’s pose estimation synchronized at a rate of 20 Hz. In each scene, the instrument was manually moved in the head model. Only one instrument was used at this stage because the attachments and markers make it impossible to simultaneously insert a second instrument.

We calibrated the endoscopic camera with the method proposed by Zhang et al. [34]. A total of 5301 data pairs were prepared. Given that the contours have to be manually annotated, we used 20 images from two sequences to estimate the mounting error. The hyperparameters $\alpha$ and $d_{max}$ were set as 0.8 and 10, respectively. The mounting errors were searched in the interval $\left[-1,1\right]$ mm for each translation dimension and $\left[-1,1\right]$ degree for each orientation dimension. These are reasonable boundaries for our fabrication and mounting accuracy.

We tested with another 20 manually annotated images from the unused sequence using the 2D projection overlap metric (1). The results are summarized in Table II. The results on the test images suggest that the mounting errors decreased owing to our method. The maximum estimated mounting error was 0.6 mm and 0.5 degree. Nonetheless, there was still around 2-3 mm of error between the projections. We believe the remaining error has two sources. First, ternary search cannot always obtain the optimal value and in our case depends on hyperparameters. Second, the visual tracker itself has an associated accuracy and our best estimates can only be as good as the tracker’s. The tracker can estimate the position of each passive marker with a 0.12 mm RMS error, that error is amplified by assembly and fabrication errors in the adapter between the marker and the tip of the instrument or endoscope.

We used translation error, centerline angle error, mean average precision (mAP) and the detected rate as the performance metrics. The centerline angle error is obtained from the angle between the estimate instrument shaft centerline and the real one.

In this section, the effect of each DR component is evaluated in our target application. We trained the model proposed in [6] for 120,000 iterations with a batch size of 20 and the one-cycle policy learning rate [36] wherein the learning rate begun at 10^-5, increased to 10^-4, and subsequently decreased to 5$\times$10^-7.

The pose map has seven channels corresponding to $x$, $y$, and $z$ translations; roll, pitch, and yaw rotations, respectively, and gripper angle. The pose loss was weighted with empirical factors of 1, 1, 1.5, 2, 2, 0.1, and 0.1.

We used the following data augmentation strategies in all cases: random contrast, random hue, random saturation, and additive noise augmentation [6]. We did not use real images for training and instead concentrate on the effects of each synthetic dataset.

The results of this ablation study are summarized in Fig. 8. Interestingly, some DR strategies deteriorated the generalizability of the model to real images. For instance, the randomization of the instrument’s glossiness strongly deteriorated the generalizability. We infer that this means that metallic reflection is important for detecting and estimating the pose. Non-contextual images used as background textures improved all metrics. In addition, the proposed normal-map-based random scratch generation improved all metrics, although further randomization with non-contextual convex and concave deteriorated the performance.

For this evaluation, we selected the DR components that, on average, improved the metrics in the ablation study. We compared our strategy with the cut-and-paste method [12], without Poisson Blending.

The evaluation results were Table IV. The proposed renderer had better performance compared to the cut-and-paste method. Ray-traced reflections are important for pose estimation in our context.

Furthermore, by adding the optimized DR components, the detected rate and translation were improved although rotation error increased. DR improves the object detection rate. In terms of pose estimation, DR was only effective if the correct DR components were used.

The performance of the proposed model was also evaluated. For this purpose, 10,000 photorealistic synthetic images without DR were used as test data. In addition, 90,000 synthetic images with all DR components are used as the training data. Other training settings were the same as that mentioned in Section IV-D.

The implementation details of the pose estimation model were as follows. ResNet34 [37] was used as the first encoder. The output of the decoder had 16 channels, and it was concatenated with the original input image to input the second encoder. The second encoder was based on ResNet18 but extended with two additional layers to detect instruments when they make most of the image, for instance, when they are close to the endoscope. Each layer comprised two residual blocks with a channel size of 512. The loss function is based on multi-task loss defined in [6] and binary cross-entropy loss for mask guidance was added, with a factor of 4.0. To compare with [6] in which only the shaft was estimated, the orientation angle error calculation does not consider the $z-$axis rotation.In this experiment, we also add a rigid transformation augmentation [31], to allow rotation augmentation of the images paired with a suitable rotation of the paired instrument pose. We rotated in the range of [-15^∘, 15^∘] with 80% probability. The inference speed was measured with Python 3.6 using Pytorch 1.4 and CUDA 10.0 on Ubuntu 18.04 with a Titan RTX graphics card.

The results are summarized in Table V. The proposed Double Encoder model is fast and accurate. By including mask guidance, pose estimation was further improved. Moreover, the rigid geometric augmentation improved results. The 90,000 synthetic images used can be considered a small number given that 8.6$\times 10^{9}$ images are needed if a training dataset is created at intervals of 1 mm and 2^∘. Therefore, the proposed method could estimate the pose with a relatively coarse interpolation of the training data.

Finally, we report the best performance with the proposed PBR, optimized DR, and proposed pose estimation model in the pose estimation of real images. In total, 90,000 synthetic images with optimized DR were used as the training data. Moreover, an additional training method using a small number of pose-free real images [6] was added. The model trained with synthetic data was additionally trained with the same synthetic data and pose-free 50 real images. The model was additionally trained for 10,000 iterations with each batch comprising 2 real images and 18 synthetic images. The model was trained from real images except with the pose loss. To compare with [6] in which only the shaft was estimated, the orientation angle error calculation does not consider the $z-$axis rotation.

The results are summarized in Table V. The pose estimation improved on synthetic images and real images. In qualitative terms, the 2D projections of the poses estimated using the proposed method generally seemed to be more accurate than those of the ground truth pose as shown in Fig. 9. The actual errors on real images might be smaller than Table V, but we currently do not have a better sensor to test that hypothesis. The creation of a higher-accuracy real dataset is left for future work.