Spaces:
Build error
Build error
File size: 3,832 Bytes
49f700c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 | import os
import re
from pathlib import Path
from paddleocr import PaddleOCR
def sort_text(ocr_result, y_threshold=20):
"""
ocr_result: PaddleOCR의 결과 (list of lines)
y_threshold: 같은 줄로 간주할 Y좌표 오차 범위 (픽셀 단위)
"""
final_output = []
for page in ocr_result:
if not all(k in page for k in('rec_texts', 'rec_scores', 'dt_polys')):
continue
# 1. 데이터 묶기
lines = []
for text, score, poly in zip(page['rec_texts'], page['rec_scores'], page['dt_polys']):
# 점수 필터링
if score < 0.8: continue
# [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
y = poly[0][1] # 상단 좌측 y값
x = poly[0][0] # 상단 좌측 x값
lines.append({'y': y, 'x': x, 'text': text, 'score': score})
if not lines: continue
# 2. Y좌표 유사도에 따른 행(Row) 그룹화
# y좌표 기준으로 정렬
lines.sort(key=lambda x: x['y'])
rows = []
current_row = [lines[0]]
for i in range(1, len(lines)):
# 현재 텍스트와 이전 텍스트의 Y좌표 차이가 y_threshold 이내면 같은 줄
if abs(lines[i]['y'] - current_row[-1]['y']) < y_threshold:
current_row.append(lines[i])
else:
# 줄이 바뀌면 이전 줄을 X좌표 순으로 정렬하여 저장
current_row.sort(key=lambda x: x['x'])
rows.append(current_row)
current_row = [lines[i]]
# 마지막 줄 처리
current_row.sort(key=lambda x: x['x'])
rows.append(current_row)
# 3. 텍스트 병합
text_lines = []
for row in rows:
row_text = " ".join([item['text'] for item in row])
# # 같은 줄의 평균을 계산 후 맨 앞에 첨부
# avg_score = sum([item['score'] for item in row]) / len(row)
# text_lines.append(f"{avg_score:.4f}: {row_text}")
text_lines.append(row_text)
final_output.append("\n".join(text_lines))
return "\n".join(final_output)
def preprocess_document(file_path):
ocr = PaddleOCR(
lang='korean',
text_det_limit_side_len=2500,
enable_mkldnn=False
)
input_folder = file_path +'/raw'
output_folder = file_path + '/processed_paddleocr_v4_2500_noscore'
# text 파일 처리
for txt_file in Path(input_folder).glob('*.txt'):
output_file = Path(output_folder) / f"{txt_file.stem}.txt"
# 파일이 이미 존재하면 건너뛰기
if output_file.exists():
print(f'Skip (already exists): {output_file.name}')
continue
with open(txt_file, 'r', encoding='utf-8') as f:
text = f.read()
text = re.sub(r'\n+', '\n', text).strip()
with open(output_file, 'w', encoding='utf-8') as f:
f.write(text)
# jpg 파일 처리
for jpg_file in Path(input_folder).glob('*'):
# 확장자가 이미지, PDF가 아니면 건너뛰기
if jpg_file.suffix.lower() not in ['.jpg', '.jpeg', '.png', '.pdf']: continue
output_file = Path(output_folder) / f"{jpg_file.stem}.txt"
# 파일이 이미 존재하면 건너뛰기
if output_file.exists():
print(f'Skip (already exists): {output_file.name}')
continue
print(f'Processing file: {jpg_file}')
result = ocr.ocr(str(jpg_file))
extracted_text = sort_text(result)
print(f'Extracted text: {extracted_text}')
break
with open(output_file, 'w', encoding='utf-8') as f:
f.write(extracted_text)
print(os.system('pwd'))
preprocess_document('data') |