DeepLearning/OCR_

[KR_OCR] 한국어데이터셋(imageWIthBox)

new_challenge 2019. 4. 21. 11:23
반응형

AI HUB에서 제공한 데이터를 바탕으로 GTUtilily와 pickle 파일을 생성한 뒤,

텍스트박스가 잘 어노테이션 되어있는지 확인 하는 작업을 진행한다.

 

1. Loading Library

import matplotlib.pyplot as plt
import PIL
import pickle
import os
import cv2 
import numpy as np
from data_KSign import GTUtility

- 한국어 데이터셋으로 만든 GTUtility를 import 한다

 

 

2. Making Function

def getTextbox(data_path = 'data/K-Sign', su = 2, quality = 'high'):
    
    if quality == 'high':
        gt_util = GTUtility(os.path.join(data_path,'HighQuality'), quality='high')
        file_name = 'gt_util_ksign_high.pkl'
        pickle.dump(gt_util, open(file_name,'wb'))
    elif quality == 'mid':
        gt_util = GTUtility(os.path.join(data_path,'MidQuality'), quality='mid')
        file_name = 'gt_util_ksign_mid.pkl'
        pickle.dump(gt_util, open(file_name,'wb'))
    elif quality == 'low':
        gt_util = GTUtility(os.path.join(data_path,'LowQuality'), quality='mid')
        file_name = 'gt_util_ksign_low.pkl'
        pickle.dump(gt_util, open(file_name,'wb'))
    elif quality == 'cr':
        gt_util = GTUtility('data/K-Sign/Annotation/', quality='cr')
        file_name = 'gt_util_ksign_cr.pkl'
        pickle.dump(gt_util, open(file_name,'wb'))
    else:
        print('quality 다시 입력')
    
    with open(file_name, 'rb') as f:
        gt_util = pickle.load(f)

>> 함수를 부를 때 받는 파라미터로는 이미지 데이터가 위치한 경로 : data_path

>> 확인 하고 싶은 이미지 갯수를 su 파라미터로 보낸다.

>> 사용할 퀄리티 값 : high, mid, low, cr (현재 데이터 셋이 퀄리티에 따라 4가지로 구성되어있기 때문)

>> 해당 quality 파라미터에 따른 pickle파일이 다르게 존재한다.

 

 

3. Drawing Textbox with Image

idxs = range(su)
    
    for idx in idxs:
        img_name = gt_util.image_names[idx]
        img_path = os.path.join(gt_util.image_path, img_name)
        
        img = cv2.imread(img_path) 
        
        boxes = np.copy(gt_util.data[idx][:,:-1])
        texts = np.copy(gt_util.text[idx])
        
        bbox = boxes.reshape(-1,2)* [img.shape[1],img.shape[0]]
        
        bbox2 = bbox.reshape(-1,8)
        bbox3 = np.ndarray.tolist(bbox2) 
        
        xlist = []
        ylist = []
        for i in range(len(bbox3)):
            xpoint = bbox3[i][0::2]
            xpoint = xpoint + [xpoint[0]]

            ypoint = bbox3[i][1::2]
            ypoint = ypoint + [ypoint[0]]

            xlist.append(xpoint)
            ylist.append(ypoint)
        
        
        plt.figure(figsize=(20,20))
        for i in range(len(xlist)):
            plt.plot(xlist[i], ylist[i], linewidth=3, color='r', label = texts[i])
        plt.imshow(img)
        plt.axis('off')

- 이미지 파일을 cv2.imread()로 읽어서 3차원 배열로 받는다

- GTUtility를 사용해서 만들어진 pickle파일의 data와 text값을 새로운 변수에 담는다

- pickle 파일의 data안에는 텍스트 박스에 관한 정보가, text에는 해당 박스의 레이블에 관한 정보가 들어있다.

- 이미지 위해 텍스트 박스를 그리기 위해 박스포인트를 원래의 값으로 변경(현재는 0~1사이의 값으로 정규화)

- x 포인트 4개를 하나의 리스트에 담고, y 포인트 4개를 하나의 리스트에 담아준다 (xy순서는 꼭 동일하게 맞춰서)

- 그리고 마지막에 박스의 형태가 닫힌형태가 되어야하기 때문에 리스트의 첫번 째 값을 마지막에 다시 추가해준다.

- 그리고 plt.plot으로 xlist, ylist를 그린다.(이미지 위에 그려준다)

 

 

**이미지를 따로 저장해서 확인하고 싶다면 아래 코드를 사용**

plt.savefig(img_name)

 

4. Final code (def getTextbox)

ImgWithTextboxes_KSign.py
import matplotlib.pyplot as plt
import PIL
import pickle
import os
import cv2 
import numpy as np
from data_KSign import GTUtility

def getTextbox(data_path = 'data/K-Sign', su = 2, quality = 'high'):
    
    if quality == 'high':
        gt_util = GTUtility(os.path.join(data_path,'HighQuality'), quality='high')
        file_name = 'gt_util_ksign_high.pkl'
        pickle.dump(gt_util, open(file_name,'wb'))
    elif quality == 'mid':
        gt_util = GTUtility(os.path.join(data_path,'MidQuality'), quality='mid')
        file_name = 'gt_util_ksign_mid.pkl'
        pickle.dump(gt_util, open(file_name,'wb'))
    elif quality == 'low':
        gt_util = GTUtility(os.path.join(data_path,'LowQuality'), quality='mid')
        file_name = 'gt_util_ksign_low.pkl'
        pickle.dump(gt_util, open(file_name,'wb'))
    elif quality == 'cr':
        gt_util = GTUtility('data/K-Sign/Annotation/', quality='cr')
        file_name = 'gt_util_ksign_cr.pkl'
        pickle.dump(gt_util, open(file_name,'wb'))
    else:
        print('quality 다시 입력')
    
    with open(file_name, 'rb') as f:
        gt_util = pickle.load(f)
        
    idxs = range(su)
    
    for idx in idxs:
        img_name = gt_util.image_names[idx]
        img_path = os.path.join(gt_util.image_path, img_name)
        
        img = cv2.imread(img_path) 
        
        boxes = np.copy(gt_util.data[idx][:,:-1])
        texts = np.copy(gt_util.text[idx])
        
        bbox = boxes.reshape(-1,2)* [img.shape[1],img.shape[0]]
        
        bbox2 = bbox.reshape(-1,8)
        bbox3 = np.ndarray.tolist(bbox2) 
        
        xlist = []
        ylist = []
        for i in range(len(bbox3)):
            xpoint = bbox3[i][0::2]
            xpoint = xpoint + [xpoint[0]]

            ypoint = bbox3[i][1::2]
            ypoint = ypoint + [ypoint[0]]

            xlist.append(xpoint)
            ylist.append(ypoint)
        
        
        plt.figure(figsize=(20,20))
        for i in range(len(xlist)):
            plt.plot(xlist[i], ylist[i], linewidth=3, color='r', label = texts[i])
        plt.imshow(img)
        plt.axis('off')
        #plt.savefig(img_name)

>> 최종적으로 완성 된 파일이다.

>> 함수를 사용 할 때는 아래와 같이 사용한다.

getTextbox(data_path = 'data/K-Sign', su = 10, quality = 'cr')

>> OUTPUT 결과 확인

 

>> 위와 같이 해당 이미지 위의 텍스트 박스가 그려지는 것을 확인 할 수 있다.

 

반응형