Learn to Differ: Sim2Real Small Defection Segmentation Network

Template free methods can be roughly divided into models based on pretrained feature extractor, generative adversarial network (GAN), and convolutional autoencoder.

Methods based on pretrained feature extractors usually utilize widely adopted feature extractors as well as classifiers. Napoletano et al. [14] utilized ResNet-18 [10] trained on ImageNet [11] to obtain a clustered feature extractor to recognize defections. Andrews et al. [2] achieved the segmentation with the backbone of VGG [18] and utilized a One-Class-SVM [5] to recognize the defection. These methods are limited by the classifier and are gradually replaced by the following generative methods.

GAN based and autoencoder based methods are generative with the assumption that defections should not be generated since they are unseen in the training set. Schlegl et al. [17] proposed a GAN based manifold extractor that was trained on defection free images. It learned a representation of the manifold of the clear images and expecting it to recognize and localize the defections in a defected image. Ackay et al. [1] proposed “Ganomaly” and utilized an autoencoder as the generating part of the GAN following by a decoder to generate a positive image. By comparing the generated image to the input image pixel-wise, defections are expected to be revealed. Bergmann et al. [4] generated positive images with an autoencoder and constructed a specific SSIM loss so that the generative error is related to the pixel-wise similarity. Zhai et al. [20] introduced energy to the regularized autoencoder so the higher energy in the generated output represented the defections. Rudolph et al. [16] detected anomalies via the usage of likelihoods provided by a normalizing flow on multi-scale image features. Comparing to these generative methods, the proposed SSDS adopts the generative fashion but differs in the content of the generated output: SSDS generates the defection directly from the defected and defection-free inputs so that even unseen defections can be recognized.

Different from template free approaches, template based methods usually coupe with tiny, vital but untrained defections which can only be recognized by comparing the defected image and the defection-free template. The origin of this is image subtraction [12] by XOR logical operator, it is easy to apply but with extremely high requirements for pixel-wise alignment of the images. Considering this, Ding et al. [8] with feature matching as an improvement has been proposed since it is more robust with pixel errors. Besides feature matching of two images, Gaidhane et al. [9] utilized similarity measurement with symmetric matrix to localize the defection. Although these methods succeeded in detecting tiny untrained defections, they are heavily influenced by the context which means the generalization capability is unsatisfying.

SSDS in this paper aims at detecting those untrained, multi-scaled defections with template based inputs but in a generative way. The proposed method is able to generate defection maps between defected and defection-free images regardless of the context so that it can even be trained in simple simulated data to generalize to multiple practical applications where defection-free template is available.

Given two images $I_{1}$ and $I_{2}$ with pose transformations, a variation of the deep phase correlator [6] is utilized to estimate the the overall relative pose $\xi_{i}$ between $I_{1}$ and $I_{2}$. The rotation and scale is calculated by

	$$\displaystyle p(\xi^{\theta,s}_{i})=\mathfrak{C}(\mathfrak{L}(\mathfrak{F}(I_{2})),\mathfrak{L}(\mathfrak{F}(I_{1})))$$		(1)
	$$\displaystyle(\theta_{i},s_{i})=\mathop{\mathbb{E}}(p(\xi^{\theta,s}_{i}))$$		(2)

where $\xi^{\theta,s}_{i}$ is the rotation and scale part of $\xi_{i}$ between $I_{1}$, $\mathfrak{F}$ is the discrete Fourier Transform, $\mathfrak{L}$ is the log-polar transform, and $\mathfrak{C}$ is the phase correlation solver. Fourier Transformation $\mathfrak{F}$ transforms images into the Fourier frequency domain of which the magnitude has the property of translational insensitivity, therefore the rotation and scale are decoupled with displacements and are represented in the magnitude. Log-polar transformation $\mathfrak{L}$ transforms Cartesian coordinates into log-polar coordinates so that such rotation and scale in the magnitude of Fourier domain are remapped into displacement in the new coordinates, making it solvable with the afterward phase correlation solver. Phase correlation solver $\mathfrak{C}$ outputs a heatmap indicating the displacements of the two log-polar images, which eventually stands for the rotation $\theta_{1}$ and scale $s_{1}$ of the two input images $I_{2}$ and $I_{1}$. To make the solver differentiable, we use expectation $\mathop{\mathbb{E}}(\cdot)$ as the estimation of $\cdot$.

Then $I_{2}$ is rotated and scaled referring to $\xi^{\theta,s}_{i}$ with the result of $\bar{I}_{2}$:

$$\bar{I}_{2}=\begin{bmatrix}s_{i}\boldsymbol{R}_{\theta_{i}}&0,\\ 0&1\end{bmatrix}I_{2}.$$

(3)

With the same manner, translations $\xi^{\textbf{t}}_{i}$ between $\bar{I}_{2}$ and $I_{1}$ is calculated

	$$\displaystyle p(\xi^{\textbf{t}}_{i})=\mathfrak{C}(\bar{I}_{2},I_{1}),$$		(4)
	$$\displaystyle\boldsymbol{t}_{i}=\mathop{\mathbb{E}}(p(\xi^{\textbf{t}}_{i})).$$		(5)

One unique property of deep phase correlation is to identify whether the two input images are related in content. When and only when the two inputs of the phase correlation in the rotation estimation stage are completely irrelevant, the phase correlation map $p(\xi^{\theta,s})$ will be a one-peak distribution centering at exactly $180\degree$ [15]. With this property in mind, we will move on to the difference learning module, the difference learning.

In this module, we aim at designing a learnable layer $\mathcal{N}_{d}$ that can recognize differences between two input images $I_{s}$ and $I_{t}$ without being disrupted by the context, and this in the practical applications could be defections in PCB, fingerprints in screen manufacturing, etc.

$$\displaystyle O_{t},O_{s}=\mathcal{N}_{d}(I_{t}||I_{s}),$$

(6)

where $||$ is the concatenation, $O_{t}$ and $O_{s}$ are the desire outputs of $\mathcal{N}_{d}$ with respect to $I_{t}$ and $I_{s}$. The content of $O_{t}$ and $O_{s}$ is further supervised by the proposed “Irrelevance Loss” $\mathcal{L}_{irr}$ and by the outlier masking layer $\mathcal{N}_{m}$ introduced in Section III-C.

The “Irrelevance Loss” $\mathcal{L}_{irr}$ forces the content of $O_{t}$ and $O_{s}$ to be completely different by utilizing the similarity identifying ability of deep phase correlation introduced in Section III-A:

$$\displaystyle p(\xi^{\theta,s}_{d})=\mathfrak{C}(\mathfrak{L}(\mathfrak{F}(O_{t})),\mathfrak{L}(\mathfrak{F}(O_{s}))),$$

(7)

where $p(\xi^{\theta,s}_{d})$ is the phase correlation result on rotation and scale with respect to $O_{t}$ and $O_{s}$. By expecting $O_{t}$ and $O_{s}$ to have completetly irrelevant content, $p(\xi^{\theta,s}_{d})$ is then supervised by the Kullback–Leibler Divergence loss $KLD$ between the Gaussian Blurred one-peak distribution centering at $180\degree$:

$$\mathcal{L}_{irr}=KLD(p(\xi^{\theta,s}_{d}),\textbf{1}_{180\degree}).$$

(8)

By now the content of $O_{t}$ and $O_{s}$ is expected to be completely different from each other. Resulting from this, there are two possibilities for $O_{t}$ and $O_{s}$: i) they are totally in a mess; ii) one of them contains useful information of differences between $I_{t}$ and $I_{s}$. This leads to the next module “Outlier Masking” which also constrains the content of $O_{t}$ and $O_{s}$ to be useful.

Inspired by the problem in Section III-B, in this module, we define the content of $O_{t}$ and $O_{s}$. We start with focusing on the final defection segmentation result $O$ who is derived from the $O_{t}$ and $O_{s}$ with a learnable masking layer $\mathcal{N}_{m}$:

$$O=\mathcal{N}_{m}(O_{t}||O_{s}).$$

(9)

The result $O$ should be exactly the same as the ground truth of defection In the coordinate system of $I_{s}$, shown as “Defection Estimated” and “Defection GT” in Fig. 1 respectively. Therefore, a hard supervision based on Mean Square Error $MSE$ on $O$ and the ground truth $G$ is applied:

$$\mathcal{L}_{def}={\lVert O-G\lVert_{2}}^{2},$$

(10)

where ${\lVert\cdot\lVert_{2}}^{2}$ is the $MSE$ loss of $\cdot$. The final loss $\mathcal{L}$ should be:

$$\mathcal{L}=\mathcal{L}_{def}+\mathcal{L}_{irr}.$$

(11)

Since a masking layer $\mathcal{N}_{m}$ cannot generate defections out of nowhere, at least one of $O_{t}$ and $O_{s}$ should contain the defection features, However, with $\mathcal{L}_{irr}$ introduced (in Section III-B), $O_{t}$ and $O_{s}$ must not contain same objects so that the defection features should only be presented in one of them, then we have:

$$\exists!~{}X\in\{O_{t},O_{s}\}(X\sim O),$$

(12)

which stands for: there exists exactly one variable in the set $\{O_{t},O_{s}\}$ that is similar to the output $O$.

To this end, it is clear that one of $O_{t}$ and $O_{s}$ is a noisy defection map, but we are not concerned about who exactly it is since they are both intermediate results to be processed by $\mathcal{N}_{m}$. This explains that the difference learner learns $\mathcal{N}_{d}$ differences between the two input images and that the context is neglected. With the masking layer proposed, the content of $O_{t}$ and $O_{s}$ is defined and the segmentation is completed.

The training skills also play important roles in the methodology since we are designing a once-for-all network needless of further training. One simulated dataset is specifically designed for the SSDS’s training with randomization, shown in Fig. 3 (Simulation Training Set). One defection-free template is firstly generated with a random number of random shapes at random locations. For the defected image, comparing to the defection-free one, the same number of same shapes are generated around the same location with slight random translational disturbance for imitating practical usage, e.g. assembly error of electronic components. Then the generated defected image is rotated and scaled randomly to simulate the overall pose misalignment in practical applications. The simulated data is completely randomized so that the network will not overfit and will focus on the difference learning. Note that the shapes are partially borrowed from [3].

For all dataset above, defected images are pose transformed from the defection-free template with the constraint on translations of both $x$ and $y$, rotation changes and scale changes of the two images in the range of $[-50,50]$ pixels, $[0,\pi)$ and $[0.8,1.2]$ respectively with images shapes of $256\times 256$.

We evaluate the performance of the defection segmentation with Intersection-over-Union(IoU) AP, Recall Rate R and the Harmony Mean(MaxF1) of AP and R. In Section V, AP and MaxF1 are adopted for quantitative demonstration.

Baselines in the experiment include template free methods GANomaly[1], SSIM [4], DSEBM[20] and template based method TDD-net [8]. All the comparative methods share the same training condition and are trained in the simulated dataset and on HiDefection. They are evaluated on HiDefection as comparisons to prove the generalization capability of SSDS.

In this experiment, we evaluate SSDS on HiDetection:SIM where “Scene I” has simulated shapes on PCBs and “Scene II” has randomly placed electronic components on PCBs. We train SSDS and the rest of the baselines in simulation and test it on HiDetection:SIM. However, to prove the generalization capability, we additionally train all baselines in HiDetection:SIM except for SSDS. The qualitative result shown in Fig. 4 and the quantitative result in TABLE. I demonstrate the superior performance SSDS achieved in generalizing simulation training to practical inference and it outperforms other baselines in generalizing. Even when trained in simulation with a completely different context, SSDS could still be on par with baselines that are trained on HiDetection:SIM when comparing the performance on inferring defections of PCBs. When inferred in a different context from training, template free and generative baselines (GANomaly, SSIM, DSEBM) are easily confused by the new context and inevitably miss important defections, and template based baseline TDD-net is misled by the assembly error of the components so that its output is blurry.

In the HiDetection:SIM where defections are simulated and are similar to what they are in the simulated training, baselines’ generalizing ability is still worthy of segmentation. However, when defections are from the real world, their performances start to drop, and this leads to Section V-B.

In this experiment, we evaluate SSDS on HiDefection:REAL where “Scene III” and “Scene IV” are both real-world particles and PCBs with lighting changes. Different from what has been conducted in Section V-A, all methods including SSDS have not been trained on HiDefection:REAL so that the generalization is further studied. The qualitative result is shown in Fig. 5 and the quantitative result is shown in TABLE. II. The result indicates that with defections from the real world involved which have not been seen in any of the training data, both the comparing methods trained in simulation and trained in HiDefection:SIM fail to present successful predictions of the defection. In contrast, SSDS trained in simulation is still a competent approach to recognize defection, even when light changes are involved in one pair of images. This proves the generalizing ability of SSDS and that the method can be directly adopted in practical applications without any further training.

Several experiments are designed for ablation studies to reassure that each module involved plays its own important role. Firstly, we validate the role of the difference learning layer by deprecating this layer and the corresponding loss in the training stage(denotes as “w/o $\mathcal{L}_{irr}$”). In such study, “w/o $\mathcal{L}_{irr}$” is trained both in simulation and on HiDefection and evaluated in HiDefection. The qualitative result (Fig. 6) and the quantitative result (TABLE. III) shows that while “w/o $\mathcal{L}_{irr}$” trained on HiDefection still gives a satisfying result, it fails in defection segmentation when it is trained in the simulation. This proves that, without the difference learning layer, the method will still work when the training and inferring share the same context, but will fail when they are different. It also proves that the difference learning layer is a vital component for the context insensitive behavior.

We further study the importance of the masking layer by deprecating such layer and calculate the defection MSE loss $\mathcal{L}_{def}$ directly between $O_{s}$ and $G$ (denotes as “w/o Mask”). The qualitative result (Fig. 6) and the quantitative result (TABLE. III) shows that “w/o Mask” trained on both HiDefection and simulation is able to provide acceptable results in but with more noise generated from the assembly error on components. This proves that the masking layer efficiently eliminates the unfavorable noise due to pixel error.

We visualize the intermediate results in the whole method to prove our interpretation of the network, shown in Fig. 7. The first case (upper case in Fig. 7) is conducted on a set of PCBs images shot in different lighting conditions, and the second case (lower case in Fig. 7) is where the shapes of silicon grease varies from each other but should not be counted as defection. We trained the network twice in simulation resulting in two independent models. The first model have the “Intermediate 2” generated with defections while the second model have it in “Intermediate 1”. It proves that one of the two outputs of the “Difference Learner” contains specific defections with noise introduced by pixel-wise error of the two images and the “Outlier Mask” is able to filter the noise out.

While it is satisfying to prove that SSDS is able to deal with lighting changes, its potential of handling images with false defections is rewarding. Note that in the second case, when defections are presented in “Intermediate 1”, the differences in silicon grease is presented in “Intermediate 2” and is filtered out after the “Outlier Mask”.