TextRearrange/DBSCAN.py

import os
import glob
import time
import numpy as np
from os import path
import pandas as pd
import seaborn as sns
import matplotlib.pyplot  as plt
from sklearn.cluster import DBSCAN
from sklearn.preprocessing import StandardScaler
import arrange as DbscanArrange
import directories as Dir

"""

이전 프로세스:
    image detection을 통해 handwritten text image에서
    각각의 word image를 추출했으나,
    순서를 알지 못하게 단어만 확인이 되어
    원래 글의 의미를 잃어버리는 문제가 있었음.
    다만 추출 시에 단어의 좌표값을 포함시킬 수 있었기에
    각 단어의 이름에 x,y좌표값과 w,h 정보를 포함하여
    단어 이미지 파일을 생성


현재 프로세스:
    yolo model을 통해 crop된 이미지들은
    파일 이름에 crop되기 전 raw image에서의 좌표 정보를 포함합니다.
    x89y147w199h184.jpg
    x10y148w157h184.jpg
    x28y149w108h180.jpg

    raw image파일에서 같은 line에 있던 word들은
    가까운 y값을 갖고있을 것이므로,
    밀도기반 군집화 알고리즘을 사용하여
    같은 line에 있을 것이라고 예상하는 word값들은
    동일 클러스터에 속하게 됨.
    이렇게 동일 클러스터에 속하는 y값들의 평균을 구하여
    word의 새로운 y값으로 label.
    단어 이미지 파일 rename할 때, y값이 먼저 오게 한다.
    yyyy_xxxx.jpg

    y값에 대한 정렬이 되고
    x값에 대해 오름차순으로 정렬됨.
    0148_0010.jpg
    0148_0028.jpg
    0148_0089.jpg
    word file의 순서가 자동적으로 원래 text의 의미를 가지게 된다.

코드 설명:
    extract_text_from_filename(),
    get_folder_contents_with_text() 함수를 이용해
    파일 이름에서 x값과 y값을 추출.
    StandardScaler로 표준화.
    Get Clustered y values using DBSCAN.
    rename_file()을 이용해서 rename image files with y-clustered values.

"""
# 사용 예시
# Get cropped word images
# folder_path = "C:/Users/ban/TEXTAI/yolov5/runs/detect/yujin_paper/crops/word" -> cropped word folder

x_texts, y_texts, name_jpg = DbscanArrange.get_folder_contents_with_text(Dir.folder_path)
file_name = pd.DataFrame(name_jpg)
file_name.columns=['file_name']

#더미 리스트 생성, 넘파이 변환 후 reshape (-1, 1)
zero_list = [0 for _ in range(len(y_texts))] #2차원 맞추기 위한 zero 리스트 
zero_list = np.array([zero_list]).reshape(-1, 1)
y_text = np.array([y_texts]).reshape(-1, 1)
#print('y_text\n', y_text)
#print('zero_list\n', zero_list)


#####################################
# 표준화
scalerX = StandardScaler() # 스케일 함수 가져와서
scalerX.fit(y_text.data) # 스케일
std_y_text = scalerX.transform(y_text.data).reshape(-1, 1) #스케일 완료

feature = pd.DataFrame(std_y_text) # 스케일된 넘파이 자료형을 데이터프레임 형태로 변환
feature.columns=['feature']

data_list = [std_y_text,zero_list] # 표준화된 데이터와 더미 삽입
data = pd.DataFrame(data_list[0])  #
labels = pd.DataFrame(data_list[1])

# 데이터 컬럼명 설정
labels.columns=['labels']
data.columns=['y']
# 두 데이터열 병합
datadf = pd.concat([data,labels],axis=1)


###################################
# create model and prediction
model = DBSCAN(eps=0.04,min_samples=2)
predict = pd.DataFrame(model.fit_predict(feature))
predict.columns=['predict']

# file_name, feature, predict 병합
r = pd.concat([file_name,feature,predict],axis=1)

#r.to_csv('C:\\Users\\ban\\Desktop\\predict_final.csv')
###########################################################



###########################################################
r = r.sort_values(by=['predict'])
#print(type(set(r['predict']))[0])
predict_list = list(set(r['predict'])) # predict 의 숫자들을 수집
unknown_words = []
same_line = []
whole_word_map = []
final_result = {'y_mean' : [],
                'x_value' : [],
                'file_name' : []}
df_final_result = pd.DataFrame(final_result)
whole_word_map_df = pd.DataFrame(whole_word_map)

for _,line in enumerate(predict_list): # 하나씩 추출 -1, 0, 1, 2, ...진짜 힘들다
    if line >= 0: # predict 값이 0 이상이면,
        y_list = [] # 클러스팅된 y값들의 평균을 구하기 위한 리스트
        print(type(r['predict']))
        same_line = r[r['predict'] == line] # r 데이터프레임에서 X인 predict를 가지고 있는 열을 가져와서
        file_num = 0
        y_mean_column = []
        total_word_map = []
        total_word_map_df = pd.DataFrame(total_word_map)
        for filename in same_line['file_name']: # 같은 클러스터에서 파일 하나씩 뽑아서
            x_data, y_data = DbscanArrange.extract_text_from_filename(filename) # 해당 파일의 x, y값을 뽑아서
            y_list.append(int(y_data)) # y값 평균을 위한 y값 리스트에 삽입
            #x_file = {x_data:filename} # key값에 x값, value값에 file_name, n_line 데이터프레임에 삽입
            file_num += 1
            word_map = {'x_value' : [int(x_data)],
                        'file_name' : [filename]}

            word_map_df = pd.DataFrame(word_map)
            total_word_map_df = pd.concat([total_word_map_df, word_map_df])
            total_word_map_df = total_word_map_df.sort_values(by=['x_value'])

        y_mean = int(np.mean(y_list)) # 한 줄에 대한 평균값 얻음 #y_mean 열에 삽입
        total_word_map_df['y_mean'] = y_mean # 추가된 행 수만큼 y_mean 열 추가

    else:
        total_word_map = []
        same_line = r[r['predict'] == line]
        for filename in same_line['file_name']:
            x_data, y_data = DbscanArrange.extract_text_from_filename(filename)
            unknown_words = [r['predict']==line]
            y_mean = int(y_data)
            word_map = {'y_mean' : [y_mean],
                        'x_value' : [x_data],
                        'file_name' : [filename]}
            word_map_df = pd.DataFrame(word_map)
            total_word_map_df = pd.concat([total_word_map_df, word_map_df])

    whole_word_map_df = pd.concat([whole_word_map_df,total_word_map_df])


file_name = list(whole_word_map_df['file_name'])
x_value = list(whole_word_map_df['x_value'])
y_mean = list(whole_word_map_df['y_mean'])
whole_list = [file_name,x_value,y_mean]


cnum = 0
timestr = time.strftime("%Y%m%d%H%M%S")

#전체 데이터 셋에 대해서 하나씩 추출하여 패딩 후 이름 변환
for i in range(len(file_name)): #_, x_value, file_name, y_mean
    old_path = str(Dir.folder_path) +"/"+ str(file_name[i])
    new_path = str(Dir.folder_path) +"/"+ str(y_mean[i]).zfill(4) +"_"+ str(x_value[i]).zfill(4) + ".jpg"

    DbscanArrange.rename_file(old_path, new_path)

if os.listdir(Dir.folder_path) == True:
    folder_contents = os.listdir(Dir.folder_path)
    print('UNKNOWN WORDS: \n',folder_contents)