DRÆM – A discriminatively trained reconstruction embedding for surface anomaly detection 코드 리뷰

서론

DRÆM – A discriminatively trained reconstruction embedding for surface anomaly detection 모델 코드에 대해 리뷰한다.

코드 리뷰

Dependency

아래는 코드는 필요한 모듈을 import한다. learn_rate값을 가져오는 get_lr함수와 weight_init이라는 함수를 통해 conv layer와 batchNorm layer의 가중치를 초기화 한다.

train_DRAEM (1)

Tensorboard

visualizer = TensorboardVisualizer(log_dir=os.path.join(args.log_path, run_name+"/"))

TensorboardVisualizer : 학습 과정을 시각화를 위해 Tensorboard 객체를 생성한다

def visualize_image_batch(self,image_batch,n_iter,image_name='Image_batch'): 이미지 배치를 시각화하는 메서드를 정의한다. make_grid 함수를 통해 이미지 배치를 그리드 형태로 변환하고, self.writer.add_image를 통해 TensorBoard에 이미지를 추가한다.
def plot_loss(self, loss_val, n_iter, loss_name='loss'): loss를 시각화하는 메서드를 정의한다. self.writer.add_scalar를 통해 TensorBoard에 스칼라 값을 추가합니다.

Reconstructive subnetwork

input을 3채널로 받고 base_width를 3으로 시작한다.

model = ReconstructiveSubNetwork(in_channels=3, out_channels=3)

ReconstructiveSubNetwork : ReconstructiveSubNetwork객체를 만들어 model에 저장한다.

ReconstructiveSubNetwork는 Encoder와 Decoder형태로 구성되어 있다.

Reconstructive Encoder

Reconstructive의 Encoder를 보면 아래와 같이 구성되어 있는데, Conv2d → Batch →ReLU→MaxPool2d 형태로 layer가 반복적으로 쌓여져 있으며 CNN layer의 output dimension이 1→2→4→8 순으로 증가한다.

Reconstructive Decoder

그 다음은 Decoder에 대한 예시인데 이미지 크기를 키우기 위해서 nn.Upsample을 사용한다.
이 layer의 scale_factor는 input image의 height, width를 각 scale factor만큼 곱해진다.
예를들어 이미지의 크기가 (1, 3, 224, 224)이면 upsample 적용하면 (1, 3, 448, 448)이 된다.
여기서 mode=bilinear는 interpolation 방법이고, align_corner=True는 가장 끝단의 픽셀을 유지하게끔 하겠다는 의미다. 해당 layer에 대해 자세한 것은 [torch] Upsampling 에 적어놓았다.
up형태의 변수에서는 upsampling이 적용되고 db형태의 경우 Conv → Batch → ReLU 를 2번 반복하여 하나의 block을 구성한다.
Decoder 부분에서는 output dimension이 encoder와 반대로 8→4→2→1 순으로 convolution이 진행된다.

마지막단에서는 output channel을 3로 갖는 3x3 Conv2d layer를 적용한다.

Discriminative subnetwork

model_seg = DiscriminativeSubNetwork(in_channels=6, out_channels=2)

DiscriminativeSubNetwork : DiscriminativeSubNetwork객체를 만들어 model_seg에 저장한다.
Encoder로부터 출력 6개를 받아 decoder의 input으로 넣는다.

Discriminative encoder

Reconstructive subnetwork의 encoder와 코드와 거의 동일하며 차이점은 maxpooling과 block이 하나 더 구성돼있다. 구성은 마찬가지로 6채널로 시작해 1 → 2 → 4 → 8 → 8 형태로 간다.

여기까지(위 코드)는 Reconstructive subnetwork의 encoder와 동일하나 아래 부분에 차이가 있음

또한 Reconstructive subnetwork의 encoder는 최종네트워크만 출력하는 반면 Discriminative subnetwork는 구성된 block 6개를 전부 출력(return)한다.

Discriminative decoder

전체적 코드는 아래와 같은데 특이한 점은 111번 줄에 있는 nn.Conv2d(base_width*(8+8), base_width*8, kernel_size=3, padding=1) 부분이다. (8+8)을 해주는 이유는 Discriminative decoder의 마지막 block6의 출력이 up_b layer의 input으로 들어오고 up_b layer의 output이 Discriminative decoder의 block5와 concat되어 채널이 8개가 더 늘어나기 때문이다.

따라서 up block에서는 단순히 input하나가 들어와 2배씩 늘려가나며 db block에서는 기존 dis encoder의 출력을 concat하고 입력으로 들어가는 구조가 반복된다.

이후 마지막 fin_out layer에서는 3x3 conv를 이용해 out_channels=2로 출력한다.

이 모델은 obj_names(class)별로 Recon, Disc subnetwork를 생성한다. 만약 class가 10개이면 각 class에 대한 Recon, Disc subnetwork를 10개 만든다는 것이다.

train_DRAEM (2)

optimizer : ReconstructiveSubNetwork와 DiscriminativeSubNetwork의 optimizer Adam으로 설정한다.
scheduler : scheduler를 통해 learning rate를 조절한다.
loss_l2 : l2 Loss(MSE) , loss_ssim : SSIM, loss_focal: Focal Loss

MVTec dataset을 만들고 dataloader로 랩핑한다.

DataLoader(train)

self.image_paths : 원본 이미지 셋
resize_shape : 기본 256, 256
anomaly_source_paths : 원본 이미지와 상관 없는 합성에 쓰일 재료 이미지
augmenters : 논문에 나온 RandAugment 종류는 7개 였으나 코드에는 10개를 사용

또한 정상 이미지에 대해서는 -45, 45도 회전을 적용했다 했으나, 해당 회전은 증강 기법에 포함됨.

증강 기법에 3개를 뽑아서 randAug를 진행한다.

자세한 사항은 주석 참조.

논문에서 아래 공식에서 보여주는 beta의 위치가 상의한것 빼고는 공식과 같다.

이후 no_anomaly가 0~1 사이의 랜덤 소수값을 생성하고 50%의 확률로 그냥 원본 이미지를 return할 것인지 masking이미지를 return할지 결정한다.

다만 위 코드에서 불필요한 부분이 있는데, 138번째줄에 해당하는 augmented_image = msk * augmented_image + (1-msk)*image 코드이다.
실질적으로 코드를 일부 수정해 확인을 해보았더니 필요없는 코드로 판단된다.

확인 코드

import torch
import numpy as np
import cv2
from perlin import rand_perlin_2d_np
import matplotlib.pyplot as plt

resize_shape = (256, 256)
anomaly_source_path = "./source.jpeg"
image = cv2.imread("./original.jpeg", 1)
image = cv2.resize(image, dsize=(resize_shape[1], resize_shape[0]))
image = np.array(image).reshape((image.shape[0], image.shape[1], image.shape[2])).astype(np.float32) / 255.0

num_iterations = 3
row = 5
plt.figure(figsize=(10, num_iterations * row))

for i in range(num_iterations):
    # image 출력
    plt.subplot(num_iterations, row, i*row + 1)
    plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    plt.title("Original Image")

    perlin_scale = 6
    min_perlin_scale = 0
    anomaly_source_img = cv2.imread(anomaly_source_path)

    anomaly_img_augmented = cv2.resize(anomaly_source_img, dsize=(resize_shape[1], resize_shape[0]))

    perlin_scalex = 2 ** (torch.randint(min_perlin_scale, perlin_scale, (1,)).numpy()[0])
    perlin_scaley = 2 ** (torch.randint(min_perlin_scale, perlin_scale, (1,)).numpy()[0])

    # perlin noise 이미지 생성
    perlin_noise = rand_perlin_2d_np((resize_shape[0], resize_shape[1]), (perlin_scalex, perlin_scaley))
    
    plt.subplot(num_iterations, row, i*row + 2)
    plt.imshow(cv2.cvtColor((perlin_noise*255).astype(np.uint8), cv2.COLOR_BGR2RGB))
    plt.title("perlin_noise Image")
    threshold = 0.5

    perlin_thr = np.where(perlin_noise > threshold, np.ones_like(perlin_noise), np.zeros_like(perlin_noise))
    perlin_thr = np.expand_dims(perlin_thr, axis=2)
    
    plt.subplot(num_iterations, row, i*row + 3)
    plt.imshow(cv2.cvtColor((perlin_thr*255).astype(np.uint8), cv2.COLOR_BGR2RGB))
    plt.title("perlin_thr Image")
    threshold = 0.5

    img_thr = anomaly_img_augmented.astype(np.float32) * perlin_thr / 255.0 # 패턴을 가진 마스크만 남음

    plt.subplot(num_iterations, row, i*row + 4)
    plt.imshow(cv2.cvtColor((img_thr*255).astype(np.uint8), cv2.COLOR_BGR2RGB))
    plt.title("img_thr Image")
    
    beta = torch.rand(1).numpy()[0] * 0.8

    augmented_image = image * (1 - perlin_thr) + (1 - beta) * img_thr + beta * image * (perlin_thr)

    # augmented_image 출력
    plt.subplot(num_iterations, row, i*row + 4)
    plt.imshow(cv2.cvtColor((augmented_image*255).astype(np.uint8), cv2.COLOR_BGR2RGB))
    plt.title("augmented_image Image")

    augmented_image = augmented_image.astype(np.float32)
    msk = (perlin_thr).astype(np.float32)
    augmented_image = msk * augmented_image + (1-msk)*image

    # 여기에 augmented_image 출력
    plt.subplot(num_iterations, row, i*row + 5)
    plt.imshow(cv2.cvtColor((augmented_image*255).astype(np.uint8), cv2.COLOR_BGR2RGB))
    plt.title("final augmented_image Image")
plt.show()

아래는 직접 이미지가 생성되어가는 부분을 출력해본 결과로 final_augment_image Image 부분은 위에 설명한 불필요한 코드(if문 안에있는)를 지난 이후의 출력값이다.
따라서 if문 밖에있는(베타 곱해주는 부분)부분 이미지 augmented_image와 비교해보면 차이가 없는 것을 볼 수 있다.

train_DRAEM (3)

이 후 class 별로 학습을 수행한다.
gray_batch = sample_batched["image"].cuda() : 원본 이미지 \(I\)
aug_gray_batch = sample_batched["augmented_image"].cuda(): 합성 이미지 \(I_a\)
anomaly_mask = sample_batched["anomaly_mask"].cuda() : 마스킹 이미지 \(M_a\)
gray_rec : output of Reconstructive subnetwork
joined_in :\(I \), gray_rec을 concat
out_mask : output of Disc subnetwork
해당 코드에서 논문과의 차이점은 \(\lambda\) 를 loss에서 사용하지 않는다는 점이다.

test_DRAEM

Image level과 pixel level의 auc와 average precision 두 지표를 출력하는 부분.

obj_ap_pixel_list : 각 객체에 대해 계산된 Average Precision (AP) 값을 픽셀 수준에서 저장하는 리스트.
obj_auroc_pixel_list : 각 객체에 대해 계산된 Area Under the Receiver Operating Characteristic curve (AUROC) 값을 픽셀 수준에서 저장하는 리스트.
obj_ap_image_list : 각 객체에 대해 계산된 AP 값을 이미지 수준에서 저장하는 리스트.
obj_auroc_image_list : 각 객체에 대해 계산된 AUROC 값을 이미지 수준에서 저장하는 리스트.

anomaly_score_gt: ground truth에 해당하며 0 또는 1 값을 갖는다.
- 정상 이미지면 0, 비정상 이미지면 1
anomaly_score_prediction: 모델의 최종출력 이후 softmax → AvgPool → np.max를 통해서 나온 예측값

주석 참조.

이 코드는 저장만 하고 딱히 어떠한 출력을 하는 용도로 사용되지 않는데, 추후 모델의 출력 이미지나(image-level) 바이너리 이미지(pixel-level)를 저장하고 싶을 때 사용되는 코드로 추측한다.

out_mask_cv: softmax 이후 전처리 된 확률 값으로 추후 pixel level score를 도출
out_mask_averaged: softmax 이후 전처리 된 확률값을 Avg_pool을 적용한다.
- np.max 를 통해 image_score 도출(image-level)
flat_true_mask: true_mask_cv 를 flatten한 1D상태(ground truth of binary image)
- true_mask_cv는 Perlin mask 이미지
flat_out_mask: out_mask_cv를 flatten한 1D상태(prediction of binary image)
반복문 이후 Image level의 anomaly_score_prediction, anomaly_score_gt(0 or 1) 값을 비교

이 후 Pixel leve의 total_gt_pixel_scores(ground truth), total_pixel_scores(flat_out_mask)과 비교
최종 출력은 아래와 같다.