Holistically-Nested Edge Detection (HED) - 윤곽선 검출

Holistically-Nested Edge Detection

이 글에서는 HED (Holistically-Nested Edge Detection) 을 이용하여 윤곽선을 검출하는 방법에 대해 소개하겠습니다. 이 알고리즘은 Saining Xie에 의해 소개 되었습니다.  HED는 이미지에서 물체의 경계 또는 객체 경계를 검출 할 수 있는 심층 신경망이라 소개하고 있습니다.

디렉토리 구조는 다음과 같이 구성했습니다. 소스코드는 Github에 있습니다.

├── detectors
│ ├── hed
│ │ ├── weights
│ │ │ └── hed_pretrained_bsds.caffemodel
│ │ ├── __init__.py
│ │ └── deploy.prototxt
│ └── __init__.py
├── lib
│ └── utils.py
├── asset
│ ├── images
│ │ ├── test_01.jpg
│ │ ├── test_02.jpg
│ │ ├── test_03.jpg
│ └ └── test_04.jpg
└── HED_detection.ipynb

---

먼저 HED Caffe 모델 파일 다운로드가 필요합니다.

1. Import Packages

import time
import torch
from .models import resnet101
import cv2
import numpy as np
import datetime

2. Class definition

제가 Class로 정의하는 이유는 첫째는 사용하기 편하게 하기 위해서고, 둘째는 Output의 모습을 맞추어 비교하기 위해서 입니다. Class로 정의하지 않고 간단하게 사용 할 수도 있습니다.

PATH_WEIGHT = './detectors/rcf/weights/only-final-lr-0.01-iter-130000.pth'

class HED():

    def __init__(self):
        tstamp = time.time()
        print('[RCF] loading...')
        self.net = cv2.dnn.readNetFromCaffe(PROTO_TXT_PATH, MODEL_PATH)
        cv2.dnn_registerLayer('Crop', CropLayer)
        print('[RCF] finished loading (%.4f sec)' % (time.time() - tstamp))

    def detect_edge(self, img, width=256, height=256):
        org_width = img.shape[1]
        org_height = img.shape[0]

        start_time = datetime.datetime.now()
        print('시작시간 : {}'.format(start_time))
        img = cv2.resize(img, (width, height))
        inp = cv2.dnn.blobFromImage(img, scalefactor=1.0, size=(width, height),
                                    mean=(104.00698793, 116.66876762, 122.67891434),
                                    swapRB=False, crop=False)
        self.net.setInput(inp)
        out = self.net.forward()

        out = out[0, 0]
        out = cv2.resize(out, (org_width, org_height))

        out = 255 * out
        out = out.astype(np.uint8)

        end_time = datetime.datetime.now()
        print('종료시간 : {}'.format(end_time))

        time_delta = end_time - start_time
        print('수행시간 : {} 초'.format(time_delta.seconds) + "\n")

        return out

class CropLayer(object):
    def __init__(self, params, blobs):
        self.xstart = 0
        self.xend = 0
        self.ystart = 0
        self.yend = 0

    # Our layer receives two inputs. We need to crop the first input blob
    # to match a shape of the second one (keeping batch size and number of channels)
    def getMemoryShapes(self, inputs):
        inputShape, targetShape = inputs[0], inputs[1]
        batchSize, numChannels = inputShape[0], inputShape[1]
        height, width = targetShape[2], targetShape[3]

        self.ystart = int((inputShape[2] - targetShape[2]) / 2)
        self.xstart = int((inputShape[3] - targetShape[3]) / 2)
        self.yend = self.ystart + height
        self.xend = self.xstart + width

        return [[batchSize, numChannels, height, width]]

    def forward(self, inputs):
        return [inputs[0][:,:,self.ystart:self.yend,self.xstart:self.xend]]

3 Test

Colab 또는 Jupyter Notebook에서 이미지를 확인하기 위한 Function과 contours 찾고 그리는 Function은 utils.py 파일로 정의했습니다.

import cv2
from detectors import HED
from imutils.perspective import four_point_transform
from lib.utils import plt_imshow, find_contours, draw_contours

detection_model = HED()

[RCF] loading...
[RCF] finished loading (0.1205 sec)

테스트를 진행 할 이미지를 Load합니다.

IMAGE_PATH = 'asset/images/test_04.jpg'
img = cv2.imread(IMAGE_PATH)
plt_imshow(["Original"], [img])

test image

edge_img = detection_model.detect_edge(img, 500, 500)
plt_imshow(["img", "edge"], [img, edge_img])

오른쪽 이미지는 Holistically-Nested Edge Detection을 적용한 최종 이미지 입니다.

시작시간 : 2022-02-08 17:51:24.278781
종료시간 : 2022-02-08 17:51:25.013391
수행시간 : 0 초

edge image

이전 글에서 소개했던 스캐너 이미지로 변환하기 위해서  가장 큰 윤곽선을 찾아 원근 변환을 적용합니다.

receiptCnt = find_contours(edge_img, 255)
draw_contours(img, receiptCnt)

outline

receipt = four_point_transform(img, receiptCnt.reshape(4, 2))
plt_imshow("Receipt Transform", receipt)

result