Deep Learning - 5

Unet 정리

먼저 labelme를 이용해서 데이터를 라벨링 시킨다.

git clone https://github.com/wkentaro/labelme

labelme > examples > semantic_segmentaiton 폴더로 접근해준다.

labelme2voc.py를 실행시켜 준다.

python labelme2voc.py data_annotated --labels labels.txt d6

data_annotated는 라벨링한 파일(.json) 및 이미지(jpeg) 폴더이고 —labels labels.txt 는 라벨링한 파일에 어떤 클래스가 존재하는 지 적어준 것이다. 차량 손상에는 Scratch, Dent, Break, Stamped, Rust가 있고, 배경 background와 ignore을 추가한 총 7개 클래스가 있다.

d6 는 labelme2voc.py를 실행 한 결과값을 넣어줄 폴더이고 d6에 있는 데이터 중에 JPEGImages에 있는 것들과 segmentationClassPNG 폴더에 있는 파일들의 이름이 같아야하며, JPEG이미지들을 모두 PNG로 바꾸어 주어야 한다. 필자는 fix라는 폴더를 미리 만들어주고 아래의 코드를 실행했다.

import cv2
import numpy as np
import os
from tqdm import tqdm

Data = 'D:\labelme\examples\semantic_segmentation\d6\JPEGImages'
SavePath = os.path.join(Data, 'fix')
for i in tqdm(os.listdir(Data)):
    rgb_image = cv2.imread('{}'.format(os.path.join(Data,i)))
    cv2.imwrite(os.path.join(SavePath, '{}'.format(i[:-4]) + '.png'), rgb_image)

segmentationClassPNG 폴더에 있는 파일이 그냥 검정색으로 보이지만, 실제로 이것은 클래스를 나누어서 저장한 것이다. 나중에 opencv를 이용해서 이미지를 확인해보면, 클래스 지정이 안된 부분(background)부분은 0으로 되어있고, 나머지는 각 클래스마다 지정된 숫자로 채워져 있다. 아래에 전부 검정색으로 보이지만, 진짜 전부 같은 색이 아님에 유의하자.

그냥 검정이 아니다

image-segmentation-keras을 git clone한다.

git clone https://github.com/divamgupta/image-segmentation-keras

image-segmentation-keras folder에 dataset1이라는 폴더를 만들고 아래에 하위 폴더를 다음과 같이 작성한다.

annotations_prepped ~ 폴더에는 검정색 이미지( segmenation이미지) 를 넣어준다. images_prepped_ ~ 에는 원본 이미지 (png파일) 을 넣어준다. 그리고 annotations_prepped~ 폴더와 images_prepped 폴더를 압축해 드라이브에 올린다.

훈련!!!!

#image-segmentation을 불러오기
!git clone https://github.com/divamgupta/image-segmentation-keras

#gdrive 마운트
from google.colab import drive
drive.mount('/gdrive')
%cd /gdrive

#위에서 만든 dataset1폴더 압축을 한 뒤에 gdrive에 업로드요망
!unzip '/gdrive/My Drive/dataset1.zip' -d /tmp

import sys
import os
sys.path.append(r'/content/image-segmentation-keras')

#텐서 플로를 이 버전으로 사용하지 않으면 정확도(acc)가 급감한다.
%tensorflow_version 1.15.0

#resnet50을 기반으로 한 u_net의 모델을 불러온다
from keras_segmentation.models.unet import resnet50_unet

model = resnet50_unet(n_classes=6,  input_height=416, input_width=608  )

#모델 훈련시키기, train할 이미지와, 그에 해당하는 annotation을 찾아 링크
model.train(
    train_images =  "/tmp/dataset1/images_prepped_train/",
    train_annotations = "/tmp/dataset1/annotations_prepped_train/",
    checkpoints_path = "/tmp/vgg_unet" , epochs=6
)


#gdrive에 segmentation처리한 파일 저장
for i in os.listdir('/tmp/dataset1/images_prepped_test'):
  out = model.predict_segmentation(
      inp="/tmp/dataset1/images_prepped_test/{}".format(i),
      out_fname="/gdrive/My Drive/output_epoch5/out_{}".format(i)
  )

#마지막 이미지 1개 출력
import matplotlib.pyplot as plt
plt.imshow(out)

# evaluating the model 
print(model.evaluate_segmentation( inp_images_dir="/tmp/dataset1/images_prepped_test/"  , annotations_dir="/tmp/dataset1/annotations_prepped_test/" ) )

여기서 segmentation된 이미지는 다음과 같은 이미지 입니다.

이 이미지를 원래의 이미지와 합쳐야 합니다. opencv를 통해 합치면 

import cv2
import os
import sys
import numpy as np
from tqdm import tqdm

class_image = [(197, 215, 20),(132, 248, 207),(155, 244, 183),(111, 71, 144),(71, 48, 128),(75, 158, 50)]
ROOT_PATH = 'D:\image-segmentation-keras\dataset1'

'''_background_
Scratch
Dent
Break
Stamped
Rust'''

#함수로 하셔도 무방할 듯 합니다.

for i in tqdm(os.listdir('D:\image-segmentation-keras\dataset1\images_prepped_test')):
    rgb_image = cv2.imread(os.path.join(ROOT_PATH,'images_prepped_test', i))
    seg_image = cv2.imread(os.path.join(ROOT_PATH, 'output_epoch3', 'out_' + i))
    seg_image[seg_image == list(class_image[0])]  = 0 #배경은 효과없음
    if len(seg_image[seg_image == list(class_image[1])]) != 0: 
        a = seg_image[seg_image == list(class_image[1])]   #클래스 1 = 파란색
        a[a == class_image[1][0]] = 255
        a[a == class_image[1][1]] = 0
        a[a == class_image[1][2]] = 0
    if len(seg_image[seg_image == list(class_image[2])]) != 0:
        #seg_image[seg_image == list(class_image[2])]  = 200 #클래스 2 = 초록색
        a = seg_image[seg_image == list(class_image[2])]
        a[a == class_image[2][0]] = 0
        a[a == class_image[2][1]] = 255
        a[a == class_image[2][2]] = 0
    if len(seg_image[seg_image == list(class_image[3])]) != 0: 
        #seg_image[seg_image == list(class_image[3])]  = 145 #클래스 3 = 빨간색
        a = seg_image[seg_image == list(class_image[3])]
        a[a == class_image[3][0]] = 0
        a[a == class_image[3][1]] = 0
        a[a == class_image[3][2]] = 255
    if len(seg_image[seg_image == list(class_image[4])]) != 0:
        #seg_image[seg_image == list(class_image[4])]  = 90 #클래스 4 = 청록색
        a = seg_image[seg_image == list(class_image[4])]
        a[a == class_image[4][0]] = 255
        a[a == class_image[4][1]] = 255
        a[a == class_image[4][2]] = 0
    if len(seg_image[seg_image == list(class_image[5])]) != 0:
        #seg_image[seg_image == list(class_image[5])]  = 35 #클래스 4 = 자홍색
        a = seg_image[seg_image == list(class_image[5])]
        a[a == class_image[5][0]] = 255
        a[a == class_image[5][1]] = 0
        a[a == class_image[5][2]] = 255
    img = rgb_image + seg_image
    cv2.imwrite(os.path.join(ROOT_PATH, 'result', '{}'.format(i)), img)

파란색 = Scratch, 초록색 = Dent, 빨간색 = Break, 청록색 = Stamped, 자홍색 = Rust 이다. Scratch와 Dent는 인식이 상대적으로 잘 되지만 오탐율이 높다. 또한 나머지 클래스는 표본이 적은 탓인지 인식을 잘 하지 못한다.