딥러닝으로 월리찾기

동기

Object detection을 재밌게 응용하는 방법이 있을지 생각하다 “월리를 찾아라”에 응용하면 재미있을 것 같다는 생각으로 만들어 보게 되었습니다.

개요

모델로 Faster RCNN 형식을 사용합니다. Pascal Voc와 COCO dataset에서 실험한 논문의 값을 기본으로 Wally dataset에 맞도록 값들을 수정하였습니다. 논문에 나와있는 값으로 Faster RCNN을 만드는 과정은 이전 포스팅 “TorchVision을 사용한 Faster R-CNN(2)”를 참고해 주세요.
https://hyungjobyun.github.io/machinelearning/FasterRCNN2/

Dataset은 Kaggle에서 찾았습니다.
https://www.kaggle.com/kairess/find-waldo
18장의 적은 이미지여서 다양한 방법으로 Image augmentation을 했습니다.

Dataset

Dataset은 numpy형식으로 저장된 파일입니다. numpy로 불러오면 (18, 1760, 2800, 3)의 형태임을 알 수 있습니다. label로 월리가 있는 위치의 픽셀에 1이 있습니다. 저는 Segmentation형식이 아닌 Bounding box를 만들기 때문에 해당 라벨을 np.where을 이용해 box로 바꿔 줍니다. box형식은 2차원 list [[x1,y1,x2,y2]]이고 x1,y1이 box의 좌측 상단, x2,y2가 box의 우측 하단 좌표입니다. DataAug class는 아래에 설명되어 있습니다.

코드

class Wally(Dataset):

  def __init__(self):
    self.DA = DataAug()
    self.transform = transforms.ToTensor()
    self.images = np.load("./waldo_dataset/imgs_uint8.npy",allow_pickle=True)
    self.labels = np.load("./waldo_dataset/labels_uint8.npy",allow_pickle=True)
    self.resize = iaa.Resize({"height": 1200, "width": 1960})

  def __len__(self):
    return self.images.shape[0]

  def __getitem__(self, idx):

    image = self.images[idx]
    label = self.labels[idx]

    y_list = np.where(label)[0] #y좌표 우선
    x_list = np.where(label)[1]

    x1 = min(x_list)
    y1 = min(y_list)
    x2 = max(x_list)
    y2 = max(y_list)

    box = [[x1,y1,x2,y2]]

    image, box = self.DA.random_aug(image, box)

    image, box = self.resize(image = image, bounding_boxes = np.array([box]))
    box = box.squeeze(0).tolist()


    targets = []
    d = {}
    d['boxes'] = torch.tensor(box,device=device)
    d['labels'] = torch.tensor([1],dtype=torch.int64,device = device)
    targets.append(d)
    
    return self.transform(image), targets

Data Augmentation

Dataset이 18장이기 때문에 data augmentation을 했습니다. imgaug라이브러리에서 flip을 사용하였고 crop은 직접 코드를 작성하였습니다. crop방식은 잘린 이미지가 원본 이미지의 1/4 보다 크고 항상 월리를 포함해야 하도록 했습니다. 코드로 표현하면 아래와 같습니다.

left = np.random.randint(0,min(box[0],image.shape[1]/2))
top = np.random.randint(0,min(box[1],image.shape[0]/2))
right = np.random.randint(max(left+(image.shape[1]/2),box[2]),image.shape[1])
bottom = np.random.randint(max(top+(image.shape[0]/2),box[3]),image.shape[0])

Training때는 랜덤으로 flip, crop, flip과 crop, 적용 안됨 중 선택되도록 합니다. 또 이미지가 crop되면 넓이와 높이가 기존의 절반까지 줄어들 수 있습니다. 따라서 이미지 크기를 맞추기 위해 넓이와 높이를 원본의 70%인 1960과 1200으로 고정합니다. 해당 내용은 Wally class에 있습니다. Dataset의 <코드> 부분을 참고하세요. 입출력 형식으로 이미지는 numpy로 입력, 출력되고 box는 2차원 list로 입력과 출력이 됩니다. 사용된 외부 라이브러리인 imgaug에 대한 자세한 사항은 아래 링크를 참고해 주세요.

https://github.com/aleju/imgaug

코드

class DataAug():
  def __init__(self):
    self.iaa_flip = iaa.Fliplr(1)
  
  def flip(self, image, box):
    image, box = self.iaa_flip(image = image, bounding_boxes = np.array([box]))
    box = box.squeeze(0).tolist()

    return image, box
  
  def crop(self, image, box):
    box = box[0]

    left = np.random.randint(0,min(box[0],image.shape[1]/2))
    top = np.random.randint(0,min(box[1],image.shape[0]/2))
    right = np.random.randint(max(left+(image.shape[1]/2),box[2]),image.shape[1])
    bottom = np.random.randint(max(top+(image.shape[0]/2),box[3]),image.shape[0])

    image = image[top:bottom+1,left:right+1,:]

    x1_new = box[0] - left
    y1_new = box[1] - top
    x2_new = box[2] - left
    y2_new = box[3] - top
    box = [[x1_new,y1_new,x2_new,y2_new]]
    return image, box
    
  def random_aug(self, image, box):
    switch = np.random.randint(0,4) #0~3, 3이면 no aug
    if switch == 0:
      image, box = self.flip(image, box)
    elif switch == 1:
      image, box = self.crop(image, box)
    elif switch == 2:
      image, box = self.crop(image, box)
      image, box = self.flip(image, box)
    
    return image, box

Model

모델은 Faster R-CNN논문과 대부분 같습니다. 논문과 차이가 있는 부분만 설명하도록 하겠습니다. 논문과 동일한 모델은 처음에 링크되어 있는 이전 포스팅을 참고하세요.
먼저 anchor box의 개수와 크기입니다. crop이 이루어진 점과 월리의 크기가 다양한 이유로 안정적인 결과를 위해 anchor box 개수를 3개에서 4개로 늘렸습니다. 또 월리의 크기가 전체 이미지에서 아주 작기 때문에 anchor box의 크기는 8, 16, 32, 64로 지정했습니다. aspect ratio는 논문과 동일합니다.

두번째로 다른 점은 roi_pooling의 resolution입니다. 논문에서는 7x7사이즈로 pooling한것에 비해 저는 9x9로 하였습니다. 실험적으로 성능이 우수하여 선택했는데 그 이유로 parameter수가 증가하기 때문이라고 생각합니다.

세번째로 rpn에서 생성되는 box의 개수를 조정했습니다. 입력으로 들어오는 이미지의 크기가 Pascal Voc나 COCO dataset에 비해 커진 만큼 생성되는 box의 개수도 많아져야 한다고 생각했습니다. 그래서 rpn_pre_nms_top_n_train = 6000에서 24000으로 바꿨습니다. 하지만 nms 이후 출력되는 box의 개수는 더 줄였습니다. 왜냐하면 작은 box가 조밀하게 모이도록 결과가 출력되는 경우가 많았기 때문입니다. 같은 이유로 nms_thresh를 0.7에서 0.5로 낮췄습니다.

네번째로 월리를 찾아라 이미지에는 많은 인물이 등장하고 서로 겹쳐 있는 경우가 많으므로 box_fg_iou_thresh를 논문의 0.5에서 0.7로 올렸습니다. 정답과 더 많이 겹치는 bounding box를 사용하기 위함입니다.

마지막으로 min_size와 max_size는 이미지의 크기에 맞도록 수정하였고 출력 class도 Pascal voc의 class 개수인 21에서 월리와 배경을 의미하는 2로 수정했습니다.

코드

backbone = torchvision.models.vgg16(pretrained=True).features[:-1]
backbone_out = 512

backbone.out_channels = backbone_out

anchor_generator = torchvision.models.detection.rpn.AnchorGenerator(sizes=((8,16,32,64),),aspect_ratios=((0.5, 1.0, 2.0),))

resolution = 9

roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'], output_size=resolution, sampling_ratio=2)

box_head = torchvision.models.detection.faster_rcnn.TwoMLPHead(in_channels= backbone_out*(resolution**2),representation_size=4096) 
box_predictor = torchvision.models.detection.faster_rcnn.FastRCNNPredictor(4096,2) #2개 class

model = torchvision.models.detection.FasterRCNN(backbone, num_classes=None,
                   min_size = 1200, max_size = 1960,
                   rpn_anchor_generator=anchor_generator,
                   rpn_pre_nms_top_n_train = 24000, rpn_pre_nms_top_n_test = 24000,
                   rpn_post_nms_top_n_train=1000, rpn_post_nms_top_n_test=500,
                   rpn_nms_thresh=0.5,rpn_fg_iou_thresh=0.7, rpn_bg_iou_thresh=0.3,
                   rpn_batch_size_per_image=256, rpn_positive_fraction=0.5,
                   box_roi_pool=roi_pooler, box_head = box_head, box_predictor = box_predictor,
                   box_score_thresh=0.05, box_nms_thresh=0.7,box_detections_per_img=300,
                   box_fg_iou_thresh=0.7, box_bg_iou_thresh=0.5,
                   box_batch_size_per_image=128, box_positive_fraction=0.25
                 )
#roi head 있으면 num_class = None으로 함

for param in model.rpn.parameters():
  torch.nn.init.normal_(param,mean = 0.0, std=0.01)

for name, param in model.roi_heads.named_parameters():
  if "bbox_pred" in name:
    torch.nn.init.normal_(param,mean = 0.0, std=0.001)
  elif "weight" in name:
    torch.nn.init.normal_(param,mean = 0.0, std=0.01)
  if "bias" in name:
    torch.nn.init.zeros_(param)

Training

Optimizer로 Adam에서 learning rate = 0.001, weight_decay = 0.0005를 사용했습니다. Learning rate는 10epoch마다 gamma=0.9로 감소시켰습니다. 전체 150Epoch를 학습시켰습니다.

결과

테스트 이미지는 학습때와 이미지 크기를 맞추기 위해 이미지 크기를 1200x1960으로 고정시켰습니다. score threshold는 대부분의 box를 후보로 두기 위해 0.1로 작게 하였습니다. 그리고 월리 주변으로 작은 box가 많이 생겨 NMS threshold를 0.001로 아주 작게 하였습니다. 따라서 살짝만 겹쳐도 가장 score가 큰 box만 남도록 하였습니다.
최종적으로는 5개의 box를 출력하여 5개중 하나에 월리가 포함되는것을 목표로 하였습니다.다만 score threshold가 0.1을 넘는것이 5개가 되지 않으면 score를 넘긴 box만 출력됩니다.

아래는 출력 결과입니다. 결과를 더 잘 보이기 위해 빨간 원을 그렸습니다.

실제로 찾아야할 월리는 파란 원에 있는 월리인데 샘플 사진만 찾아냈습니다.

하나만 출력되었고 그것이 정답인것을 알 수 있습니다.

이 그림에서는 월리를 찾아내지 못했습니다.

이 그림에서는 score를 0.1로 하면 출력되는 box가 없어 0으로 하였습니다.

결론

Faster R-CNN 모델을 응용하여 월리를 찾아라에 적용 할 수 있습니다. Data augmentation, anchor box의 개수, 모델의 parameter개수, 그리고 NMS와 foreground threshold가 중요한 요소였습니다. 더 많은 데이터를 가지고 더 정밀한 hyper parameter설정을 한다면 성능을 향상 시킬 수 있을 것이라고 생각합니다.

전체 코드:
https://github.com/HyungjoByun/Projects/blob/main/Faster%20RCNN/FindWally.ipynb

참고한 사이트

https://www.youtube.com/watch?v=wIDn83OJeK4&t=783s
https://towardsdatascience.com/how-to-find-wally-neural-network-eddbb20b0b90