CA-SpaceNet: Counterfactual Analysis for 6D Pose Estimation in Space

The network architecture is given in Fig. 2. The network’s input is two images: the first one is the raw image $I$ with satellite, and the second is the image $I_{m}$ with only the background after removing the satellite. Two DarkNet-53 networks with the same weights are adopted to perform features extraction of $I$ and $I_{m}$, respectively, resulting in $\textbf{F}\in\mathbb{R}^{C\times H\times W}$ and $\textbf{F}^{m}\in\mathbb{R}^{C\times H\times W}$. After feature extraction, three feature pyramid networks $\textbf{FPN}_{f}$, $\textbf{FPN}_{pc}$, and $\textbf{FPN}_{c}$ are constructed to perform counterfactual analysis. These feature aggregation modules with the same network structure but different weights are the core components of factual path, pseudo counterfactual path, and counterfactual path. Through counterfactual analysis, the side effect can be decoupled and removed from the total effect to obtain the final TDE, i.e., unbiased feature $\{\textbf{F}^{TDE}_{n}\}^{N}_{n=1}$ after weakening background interference, where $N$ denotes the number of layers in FPN. Finally, the unbiased feature will be sent to an anchor-based keypoint detector. A PnP solver is adopted to predict the final 6D pose of the target satellite.

The rest of this section is organized as follows: In subsections III-B, III-C, and III-D, the construction of the factual path, counterfactual path, and pseudo counterfactual path are described in detail, respectively. The network quantization method adopted in this paper is briefly reviewed in subsection III-E. The training and inference processes of the CA-SpaceNet are introduced in subsection III-F.

It can be seen from the ranking of SPEED competition that the methods based on PnP solver are much more stable than the methods of directly estimating 6D pose. Therefore, this paper adopts the strategy based on a PnP solver. In this strategy, the 6D pose estimation task is divided into two subtasks: 2D keypoint detection and PnP problem. Detailed compositions of the proposed framework are shown in Fig. 2. The factual path is the central part of the CA-SpaceNet, which refers to the structure of [4]. Hu et al. [4] explored the problem of huge changes in the depth range of space-borne objects. However, they directly adopted neural networks to extract features of the whole image without considering the impact of complex background information on 6D pose estimation tasks. In some general computer vision tasks, e.g., object detection and semantic segmentation, background interference will not cause a significant decrease in performance. While in some tasks requiring high precision, such as 6D pose estimation, the background interference will affect the accuracy of keypoint detection, which will significantly deteriorate the final performance.

The factual path is designed to simulate the phenomenon of background interference. In this path, $\textbf{FPN}_{f}$ completes feature aggregation and generates features with the target satellite and irrelevant background:

$$\mathcal{F}^{f}={\color[rgb]{83,129,53}\textbf{FPN}_{f}}(\textbf{F}),$$

(1)

where $\mathcal{F}^{f}$ denotes the factual feature set $\mathcal{F}^{f}=\{\textbf{F}^{f}_{n}\},(n=1,2,\cdots,N)$ generated by different layers of $\textbf{FPN}_{f}$. This feature is regarded as the total effect (TE) in counterfactual analysis. The total direct effect (TDE) in CA-SpaceNet is replaced with TE when analyzing the factual path separately. $\mathcal{F}^{f}$ will be directly sent to the network head to perform keypoint detection, resulting in $\{(\hat{u}_{k},\hat{v}_{k})\},(k=1,2,\cdots,K)$, where $K$ denotes the number of the corners of the satellite’s cubic body. 3D loss $\mathcal{L}_{3D}$ and object class loss $\mathcal{L}_{cls}$ are adopted in this paper. Please refer to [4] for more details about these loss functions.

The idea of counterfactual analysis is to imagine a non-existent path, that is, to study the effect under the What If scenario. In space scenes, the complex satellite-earth relationship and harsh illumination conditions will cause significant changes in the background. These factors will negatively impact the feature extraction stage and eventually lead to suboptimal results. Therefore, we imagine what features will be generated through ”what if there is no target?” in the counterfactual path. A path composed of a DarkNet-53 and an $\textbf{FPN}_{c}$ is constructed to realize this assumption. The path in red box of Fig. 2 shows more details. The input of the counterfactual path is $I_{m}$ with only background information after erasing the target through the ground-truth mask. Due to the absence of the target, the generated feature map only contains background information:

$$\mathcal{F}^{c}={\color[rgb]{192,0,0}\textbf{FPN}_{c}}(\textbf{F}^{m}),$$

(2)

where $\mathcal{F}^{c}$ denotes the counterfactual feature set $\mathcal{F}^{c}=\{\textbf{F}^{c}_{n}\},(n=1,2,\cdots,N)$.

Simplified causal graphs of CA-SpaceNet are shown in Fig. 3. A causal graph is a directed acyclic graph (DAG) that consists of nodes and directed edges. The nodes denote variables, and the directed edges denote cause-effects between nodes [7, 6]. Counterfactual analysis in CA-SpaceNet aims to disentangle the pure object features from the mixed features. In Fig. 3, the factual path (a) and counterfactual path (b) constitute the ideal counterfactual analysis (c). The total direct effect feature $\mathcal{F}$ in (c) is obtained by subtracting the side effect feature $\mathcal{F}^{c}$ from the total effect feature $\mathcal{F}^{f}$:

$$\mathcal{F}=\mathcal{F}^{f}-\mathcal{F}^{c}.$$

(3)

It should be noted that two DarkNet-53 networks in the counterfactual path and factual path share the same weights and will be frozen during training of the CA-SpaceNet. The reason for adopting this strategy is to only equip the FPN modules with the ability to distinguish the background and the target to avoid the backbone becoming a confounder.

Although the unbiased feature $\mathcal{F}$ can be directly calculated by Eq. 3 theoretically, the segmentation information of the target is not available in application. Therefore, we elaborately design a pseudo counterfactual path to imitate the counterfactual feature $\mathcal{F}^{c}$, which is colored in blue in Fig. 2(a). As its name implies, pseudo means that this path is a fake path, which aims to imitate the counterfactual path. $\textbf{FPN}_{pc}$ is the main component of the pseudo counterfactual path. It takes the factual feature F as input and generates imitation feature:

$$\mathcal{F}^{pc}={\color[rgb]{16,84,151}\textbf{FPN}_{pc}}(\textbf{F}),$$

(4)

where $\mathcal{F}^{pc}$ denotes the pseudo counterfactual feature set $\mathcal{F}^{pc}=\{\textbf{F}^{pc}_{n}\},(n=1,2,\cdots,N)$.

In order to make $\mathcal{F}^{pc}$ and $\mathcal{F}^{c}$ as similar as possible, i.e., $\textbf{F}^{pc}_{n}\approx\textbf{F}^{c}_{n}$, a smoothed L1 norm loss function $sl_{1}(\cdot,\cdot)$ is adopted to measure the discrepancy between them:

$$\mathcal{L}_{sim}=\sum_{n=1}^{N}sl_{1}(\textbf{F}^{pc}_{n},\textbf{F}^{c}_{n}).$$

(5)

With minimizing $\mathcal{L}_{sim}$, $\textbf{FPN}_{pc}$ can learn the ability to disentangle the pure background feature from mixed feature F. As shown in Fig. 3(d), the counterfactual feature $\mathcal{F}^{c}$ will be replaced by the pseudo counterfactual feature $\mathcal{F}^{pc}$ in real counterfactual analysis. The final approximate total direct effect feature is calculated by

$$\hat{\mathcal{F}}=\mathcal{F}^{f}-\mathcal{F}^{pc}.$$

(6)

This strategy skillfully solves the problem that the ground-truth mask is lacking during inference of the CA-SpaceNet.

In the space scene, special working conditions and limited computing resources put forward higher requirements for the power consumption and latency of the algorithm. This paper adopts the classical LSQ-Net [35] to quantize the proposed CA-SpaceNet to a low-bit-width model and then deploys a 3-bit convolutional layer into a PIM architecture on FPGA. As shown in Fig. 4, three quantization modes are proposed. The quantization range is gradually expanded to explore the impact of quantization on CA-SpaceNet finely. This is the first work that applies network quantization methods to the 6D pose estimation task of space-borne targets.

Fig. 5(a) demonstrates the detailed quantization and fusion process. Next, we will elaborate on a convolutional layer with the kernel size of $3\times 3$ and the input size of $H^{\prime}\times W^{\prime}$ as an example. Before convolution, the weights $\textbf{W}\in\mathbb{R}^{C_{out}\times C_{in}\times 3\times 3}$ and activations $\textbf{A}\in\mathbb{R}^{C_{in}\times H^{\prime}\times W^{\prime}}$ in FP32 are quantized into low-bit features $\hat{\textbf{W}}=\lfloor clip(\textbf{W}/s_{w})\rceil$ and $\hat{\textbf{A}}=\lfloor clip(\textbf{A}/s_{a})\rceil$ , where $clip(\cdot)$ returns values within quantization limits, and $\lfloor\cdot\rceil$ returns every element in $\hat{\textbf{W}}$ or $\hat{\textbf{A}}$ to its nearest integer. After the quantization convolution

$$\hat{\textbf{Y}}=Qconv(\hat{\textbf{W}},\hat{\textbf{A}}),$$

(7)

dequantization is performed to recover the activation through rescaling factors:

$$\textbf{Y}=\hat{\textbf{Y}}*s_{w}*s_{a},$$

(8)

where $\hat{\textbf{Y}}$ denotes the quantization feature and Y denotes the recovery feature. Then the convolutional layer and its adjacent BatchNorm layer are fused to a single operation by equivalence relation:

$$\begin{aligned} \mathbf{Y}^{bn}_{(i,:,:,:)}&=\frac{\textbf{Y}_{(i,:,:,:)}-\bm{\mu}_{i}}{\bm{\sigma}_{i}}\bm{\gamma}_{i}+\bm{\beta}_{i}\\ &=\frac{\bm{\gamma}_{i}}{\bm{\sigma}_{i}}*s_{w}*s_{a}*\hat{\textbf{Y}}_{(i,:,:,:)}-\frac{\bm{\mu}_{i}\bm{\gamma}_{i}}{\bm{\sigma}_{i}}+\bm{\beta}_{i}\end{aligned},$$

(9)

where the subscript $i$ denotes the index of the output channel $C_{out}$. $\bm{\mu}$, $\bm{\sigma}$, $\bm{\gamma}$, and $\bm{\beta}$ are four types of parameters in the BatchNorm layer. Following this, the ReLU layer and scaling step in the next layer are also absorbed by a single operation.

The overall weight memory occupied by the CA-SpaceNet is greatly reduced through quantization and layer fusion, which is suitable for PIM chips with the merits of energy efficiency and avoidance of the memory wall. In PIM architecture, the quantized network is pre-loaded into BRAM, and intermediate data accessed from/to the off-chip DRAM access is entirely eliminated during inference. This paper implements a part of the CA-SpaceNet into the PIM accelerator proposed by Jiao et al. [37], which is demonstrated in Fig. 5(b). The feasibility of the deployment is confirmed on FPGA.

As shown in Fig. 2(b), the training and inference phases are separated in CA-SpaceNet. All three paths are activated during training. While learning to detect keypoints, the CA-SpaceNet is supposed to minimize the discrepancy between features generated by $\textbf{FPN}_{pc}$ and $\textbf{FPN}_{c}$. Therefore, the network is trained with a weighted combination of the loss terms:

$$\mathcal{L}=\lambda_{3D}\mathcal{L}_{3D}+\lambda_{c}\mathcal{L}_{cls}+\lambda_{s}\mathcal{L}_{sim},$$

(10)

where the loss weights $\lambda_{3D}$, $\lambda_{c}$, and $\lambda_{s}$ are set to 1, 1, and 0.25, respectively.

It should be noted that the ground-truth mask is available during training, while unavailable during inference. The pseudo-counterfactual path is constructed to replace the function of the counterfactual path to address this issue. Therefore, the counterfactual path (i.e., $\textbf{FPN}_{c}$) will be deleted during inference.

The proposed CA-SpaceNet is built on the PyTorch [41] and implemented on a system with the Intel(R) Core(TM) i7-9700K CPU @ 3.60GHz. All methods are trained on a single Nvidia Titan GPU. For all CA-SpaceNet frameworks in this paper, the DarkNet-53 pretrained on ImageNet is chosen as the backbone. Stochastic gradient descent (SGD) optimizer is adopted for network optimization with initial learning rate 1e-3, momentum 0.9, and weight decay 1e-4. The training epoch is set to 30. Online data augmentation strategies such as random shift, scale, and rotation are performed during training. Different experimental settings are adopted because of different complexity. For the SwissCube dataset, the number of minibatch is set to 8, and the resolution is set to 512×512. For the SPEED dataset, the number of minibatch is set to 4, and the resolution is set to 960×960. An Ultra96v2 FPGA board is implemented to deploy the quantization convolutional layer of the CA-SpaceNet.

The SwissCube dataset is the largest released dataset in 6D pose estimation of space-borne targets. We choose this dataset as the main benchmark. Both quantitative and qualitative experiments are carried out on this dataset.

In order to verify that the performance improvement is brought by the counterfactual analysis strategy rather than the additional 30 training epochs, we conducted another model named WDR* w. 30-Ep.. Experimental results in Tab. II show that the additional 30 training epochs lead to over-fitting, while the additional 30 training epochs with counterfactual analysis (CA-SpaceNet) achieves better performance, which confirms the superiority of the proposed framework.

Tab. III compares our method with several top-performing solutions in the Kelvins Satellite Pose Estimation Challenge. Considering the limitation of computing resources, we didn’t adopt multiple model strategies and refinement to the final 6D pose results as done in [5]. There is no official masks of the satellite in SPEED, so we utilize the cv2.convexHull function in OpenCV to connect 11 corners to generate approximate masks. As the original paper of WDR [4] did not release the code on the SPEED dataset, we reproduce this model and obtain the results of WDR*. Due to some unknown tricks, the results reproduced by us is slightly worse than WDR [4], but in the same order of magnitude (0.018 for WDR and 0.040 for WDR*). Note that this does not prevent us from evaluating the proposed framework, because the CA-SpaceNet is developed from the WDR*. Compared with the WDR* reproduced by us, the CA-SpaceNet achieves lower error on ${\textbf{e}_{q}}+{\textbf{e}_{t}}$ (0.0385 for CA-SpaceNet and 0.04 for WDR*). The decrease of error proves that the CA-SpaceNet is robust to complex backgrounds interference in 6D pose estimation of space-borne targets.

After confirming the effectiveness of the proposed framework, we quantize the CA-SpaceNet into a low-bit-width model and deploy a part of the quantized model into a real hardware accelerator.

Through quantization, floating-point weights in the network are transformed into low-bit-width values for storage, which greatly reduce the size of the network. It is more easier for the network to be deployed to devices with limited memory. Tab. V gives an occupation summary of 8-bit and 3-bit networks under mode III setting. The size of the network is reduced by 75% and 90.63% in 8-bit and 3-bit settings, respectively. The minimal memory occupation provides the basis for the real deployment on chips.