[ENG_OCR] CRNN_cocotext 데이터 전처리 작업

Datasets : COCO/Train2014

Annotations : COCO_Text.json

 

 

https://vision.cornell.edu/se3/coco-text-2/

COCO-Text: Dataset for Text Detection and Recognition | SE(3) Computer Vision Group at Cornell Tech

1. COCO_Text.json에 대한 이해

import os
gt_path = 'data/COCO-Text/'
with open(os.path.join(gt_path, 'COCO_Text.json')) as f:
            gt_data = json.load(f)

위와 같은 방식으로 json데이터를 읽고, 어떤 형태로 구성되어 있는지 파악

 

# 가져온 json 데이터는 딕셔너리 형태

    gt_data.keys() : dict_keys(['imgToAnns', 'anns', 'imgs', 'cats', 'info'])

 

# imgToAnns

    >> len(gt_data['imgToAnns'].keys()) : 63686 (이미지 갯수)

    >> gt_data['imgToAnns'].keys() 안에는 이미지의 id가 있다. 해당 키의 value는 텍스트 박스의 id (0~N개)

이미지 id에는 텍스트박스 id가 존재

 

# imgs : 이미지 정보

    >> 이미지 아이디를 이용해 gt_data['imgs']['이미지ID']를 찾으면 안의 VALUE값이 출력되는데 이미지의 정보.

imgs 안에는 이미지에 해당하는 정보가 들어있다.  

 

# anns : 어노테이션정보(text box)

    >> 텍스트박스 아이디를 사용해 gt_data['anns']['텍스트박스ID']를 찾으면 안의 VALUE값이 출력되는데 박스의정보.

텍스트박스에 대한 다양한 정보가 담겨있다. 

 

 

 

2. GTUtility(BaseGTUility)

위의 내용을 참고하면 아래 코드내용을 이해하기 좋다.

JSON 파일을 받아서 사용하기 좋게 데이터를 정제하는 Utility라고 생각하면 된다.

 

## Input

    일단 아래의 함수를 사용할 때는 파라미터 값으로 파일의 위치, polygon과 validation에 대한 T/F값을 준다.

 

## Ouput

    data_path = gt_path : 데이터의 위치

    image_path : 이미지의 위치

    image_name[] : 이미지 파일의 이름

    img_width : 이미지 넓이

    img_height : 이미지 높이

    data[] : polygon일 경우 8개의 array 배열, polygon이 아닐 경우 box = [ x, y, x+w, y+h ] 

                 : 그리고 x값(0부터 2step)은 width로 나누고, y값(1부터 2step)은 height로 나눈다. 0~1사이의 값으로 정규화

    text[] : 텍스트박스의 라벨값

 

    >> image_name, data, text가 리스트 형태로 담아져서 나간다. 

 

data_cocotext.py

import numpy as np
import json
import os

from ssd_data import BaseGTUtility

class GTUtility(BaseGTUtility):
  
    def __init__(self, data_path, validation=False, polygon=False, only_with_label=True):
        test=False
        
        self.data_path = data_path
        gt_path = data_path
        image_path = os.path.join(data_path, 'train2014')
        self.gt_path = gt_path
        self.image_path = image_path
        self.classes = ['Background', 'Text']
        
        self.image_names = []
        self.data = []
        self.text = []
        
        with open(os.path.join(gt_path, 'COCO_Text.json')) as f:
            gt_data = json.load(f)

        for img_id in gt_data['imgToAnns'].keys(): # images
            
            ann_ids = gt_data['imgToAnns'][img_id]

            if len(ann_ids) > 0:
                img_data = gt_data['imgs'][img_id]
                set_type = img_data['set']
            
                if test:
                    if set_type != 'test':
                        continue
                elif validation:
                    if set_type != 'val':
                        continue
                else:
                    if set_type != 'train':
                        continue

                image_name = img_data['file_name']
                img_width = img_data['width']
                img_height = img_data['height']

                boxes = []
                text = []
            
                for ann_id in ann_ids: # boxes
                    ann_data = gt_data['anns'][str(ann_id)]
            
                    if polygon:
                        box = np.array(ann_data['polygon'], dtype=np.float32)
                    else:
                        x, y, w, h = np.array(ann_data['bbox'], dtype=np.float32)
                        box = np.array([x, y, x+w, y+h])
                    
                    if 'utf8_string' in ann_data.keys():
                        txt = ann_data['utf8_string']
                    else:
                        if only_with_label:
                            continue
                        else:
                            txt = ''
                    
                    boxes.append(box)
                    text.append(txt)
                
                if len(boxes) == 0:
                    continue
                
                boxes = np.asarray(boxes)
                
                boxes[:,0::2] /= img_width
                boxes[:,1::2] /= img_height
                    
                # append classes
                boxes = np.concatenate([boxes, np.ones([boxes.shape[0],1])], axis=1)
                
                self.image_names.append(image_name)
                self.data.append(boxes)
                self.text.append(text)
                
        self.init()


if __name__ == '__main__':
    gt_util = GTUtility('data/COCO-Text', validation=False, polygon=True, only_with_label=True)
    print(gt_util.data)

 

 

3. 마지막으로 데이터를 PICKLE형태로 저장

 

우선 CRNN Model을 학습시키기 위해 COCO-Text.json 파일을 pickle파일로 만들어주는 작업.

기존의 참고한 깃허브에 cocotext를 pickle로 만들어주는 코드가 없기 때문에 기존코드를 참고해서 새로 작성.

 

우선 위에서 선언한 classdls GTUtility를 import 하고 

GTUtility 객체를 train과 val 따로 생성한다.

 

make_cocoPickle.py

import pickle

from data_cocotext import GTUtility

gt_util_train = GTUtility('data/COCO-Text/', polygon=True, validation = False)
gt_util_val = GTUtility('data/COCO-Text/', polygon=True, validation = True)

file_name_train = 'gt_util_cocotext_train.pkl'
file_name_val = 'gt_util_cocotext_val.pkl'

pickle.dump(gt_util_train, open(file_name_train,'wb'))
pickle.dump(gt_util_val, open(file_name_val,'wb'))

print(gt_util_train)
print(gt_util_val)

위의 코드를 실행시켜 나온 print된 화면은 아래와 같다

 

 >>

상단의 피클은 gt_util_train에 해당

Train에 해당하는 이미지에 대한 pickle

 

 

 

 

>>

하단의 피클은 gt_train_val에 해당 

Val에 해당하는 이미지에 대한 pickle