bounding box check

[Ssunbell]

from PIL import Image, ImageDraw
from IPython.core.display import display
import pandas as pd
import json
import editdistance
from nltk.tag import pos_tag
import re
from typing import *

class FindCategory():
    def __init__(self, data_path='data/test/', result_path='./', start_idx=2582, end_idx=3874):
        self.data_path = data_path
        self.result_path = result_path
        self.dataset = self.make_dataset(start_idx, end_idx)

    def NLD(self, s1:str,s2:str) -> float:
        return editdistance.eval(s1.lower(),s2.lower()) / ((len(s1)+len(s2))/2) # normalized_levenshtein_distance

    def clean_text(self,raw_string:str) -> str:
        #텍스트에 포함되어 있는 특수 문자 제거
        text = re.sub('[-=+,#?^$.@*※~!…]','', raw_string)

        return text

    def find_noun_ngram(self,questions:str, ngram:int) -> Set[Tuple[str]]:
        if ngram == 1: # unigram일 경우
            part_of_speech = {'NN', 'NNS','NNP', 'NNPS', 'POS','RP', 'CD', 'FW', 'VBG'}
        else:
            part_of_speech = {'NN', 'NNS','NNP', 'NNPS', 'IN', 'POS','RP', 'CD', 'FW', 'VBG', 'JJR', 'JJS', 'RBR', 'RBS'}

        result = set()
        questions:List[str] = questions.split()
        ngram_questions = [questions[i:i+ngram] for i in range(len(questions) - (ngram-1))]
        for question in ngram_questions:
            tmp_storage = []
            for tag in pos_tag(question):
                if tag[1] in part_of_speech:
                    tmp_storage.append(self.clean_text(tag[0]))

            if len(tmp_storage) == ngram:
                result.add(tuple(tmp_storage))

        yield from result
    def make_dataset(self, start_idx, end_idx):
        df_raw = pd.read_csv(self.result_path + 'result.csv', index_col=0)
        df_raw = df_raw.drop('6', axis=1).drop('7', axis=1).drop('0', axis=1).reset_index().drop('index',axis=1)
        df_raw = df_raw.rename(columns={
            '1':'question_id',
            '2' : 'score',
            '3':'question',
            '4':'prediction',
            '5':'answer'}
        )
        with open(self.data_path+'test_v1.0.json', 'r') as f:
            json_data = json.load(f)
        df_test = pd.DataFrame(eval(json.dumps(json_data))['data'])

        df = df_raw.iloc[start_idx : end_idx+1]
        result = []
        for _, row_test in df.iterrows():
            row = df_test[row_test['question_id'] == df_test['questionId']]
            image_file = row['image']
            img = Image.open("data/test/"+image_file.iloc[0]).convert('RGB')
            draw = ImageDraw.Draw(img)
            ocr_file = 'data/test/ocr_results/' + row['ucsf_document_id'] + '_' + row['ucsf_document_page_no'] + '.json'
            with open(ocr_file.iloc[0], 'r') as f:
                json_data = json.load(f)
            answers = set(row_test['answer'].replace("\'", '').replace('[','').replace(']','').split(','))
            ngrams = 3
            question_words = sum([[ngram_question for ngram_question in self.find_noun_ngram(row_test['question'], ngram)] for ngram in range(1, ngrams+1)], [])
            for answer in answers:
                for data in json_data['recognitionResults']:
                    indices_blue = []
                    indices_red = []
                    for i, line in enumerate(data['lines']):
                        percent_question = 0
                        words_list = [word['text'] for word in line['words']]

                        ###### question 찾기 ######

                        for words_q in question_words:
                            for word_q in words_q:
                                check = False
                                for word in words_list:
                                    ld_score = self.NLD(word_q, word)

                                    if ld_score < 0.2:
                                        check = True
                                if check:
                                    percent_question += 1

                            if percent_question / len(line['words']) > 0.6:
                                indices_blue.append(i)

                        ###### answer 찾기 ######
                        percent_answer = 0

                        for word_a in answer.split():
                            check = False
                            for word in words_list:
                                ld_score = self.NLD(word_a, word)

                                if ld_score < 0.2:
                                    check = True
                            if check:
                                percent_answer += 1

                        if percent_answer / len(line['words']) > 0.3:
                            indices_red.append(i)

                for bbox_i in indices_blue: # 질문
                    bboxes = json_data['recognitionResults'][0]['lines'][bbox_i]['boundingBox']
                    draw.rectangle((bboxes[0],bboxes[1],bboxes[4],bboxes[5]), outline=(0,0,255), width = 3)
                for bbox_i in indices_red: # 정답
                    bboxes = json_data['recognitionResults'][0]['lines'][bbox_i]['boundingBox']
                    draw.rectangle((bboxes[0],bboxes[1],bboxes[4],bboxes[5]), outline=(255,0,0), width = 3)

            result.append({'question': row_test['question'], 'answer' : list(answers), 'img':img})

        return result

    def __getitem__(self, idx):
        data = self.dataset[idx]
        print('question :',data['question'])
        print('answer :', data['answer'])
        img = data['img']
        display(img.resize((int(img.width/2), int(img.height/2))))

파란색 박스가 질문과 관련된 단어가 있는 바운딩 박스, 빨간색이 정답이 있는 바운딩 박스