File size: 4,633 Bytes
d83150d caa85f5 d83150d caa85f5 d83150d caa85f5 d83150d caa85f5 d83150d caa85f5 d83150d caa85f5 |
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 107 108 109 110 111 112 113 |
import pytesseract
import json
import numpy as np
from PIL import Image, ImageEnhance
from transformers import VisionEncoderDecoderModel, TrOCRProcessor
hf_model = VisionEncoderDecoderModel.from_pretrained("Serovvans/trocr-prereform-orthography")
processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
def remove_bleed_through(image_path,
brightness_factor=1.5):
# Загрузка изображения
pil_image = Image.open(image_path).convert('RGB')
img = np.array(pil_image)
alpha = 1.7
beta = -130
result = alpha * img + beta
result = np.clip(result, 0, 255).astype(np.uint8)
# Преобразование в PIL Image
pil_result = Image.fromarray(result)
# 2. Повышение яркости
enhancer_brightness = ImageEnhance.Brightness(pil_result)
bright_image = enhancer_brightness.enhance(brightness_factor)
return bright_image
# Функция распознавания текста с TrOCR
def recognize_row(image):
image = image.convert("RGB")
pixel_values = processor(images=image, return_tensors="pt").pixel_values
generated_ids = hf_model.generate(pixel_values)
generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
return generated_text
# Основная функция распознавания текста с разделением на слова/фразы
def recognize_page(image_path, text_output_path=False):
# Открываем изображение
image = remove_bleed_through(image_path)
# Используем pytesseract для получения данных по каждому фрагменту
data = pytesseract.image_to_data(image, config="--psm 6", output_type=pytesseract.Output.DICT, lang='ukr+eng')
with open("rec_data.json", "w", encoding="utf-8") as json_file:
json.dump(data, json_file)
pad = int(0.0042 * image.size[1])
# Сохраняем результаты для последующего восстановления порядка
fragments = []
for i in range(len(data['text'])):
if data['conf'][i] > -1 and data['width'][i]*data['height'][i] > pad**2: # Игнорируем пустые фрагменты
x, y, w, h = data['left'][i], data['top'][i], data['width'][i], data['height'][i]
fragment_image = image.crop((x-pad, y-pad, x + w + pad, y + h + pad))
text = recognize_row(fragment_image).strip()
if data['text'][i].strip() == "\u2014":
text = "\u2014"
try:
num = int(data['text'][i].strip())
if data['conf'][i] > 85:
text = data['text'][i].strip()
except:
pass
fragments.append({
'block_num': data['block_num'][i],
'par_num': data['par_num'][i],
'line_num': data['line_num'][i],
'word_num': data['word_num'][i],
'text': text,
'image': fragment_image
})
# Сортируем фрагменты по line_num и word_num
fragments = sorted(fragments, key=lambda x: (x['block_num'], x['par_num'], x['line_num'], x['word_num']))
# Распознаем текст из каждого фрагмента
result_lines = []
current_line_num = 0
curr_block_num = 0
curr_par_num = 0
current_line = []
for fragment in fragments:
if fragment['line_num'] != current_line_num or fragment['block_num'] != curr_block_num or fragment['par_num'] != curr_par_num:
# Завершаем текущую строку и переходим к следующей
result_lines.append(" ".join(current_line))
current_line = []
current_line_num = fragment['line_num']
curr_block_num = fragment['block_num']
curr_par_num = fragment['par_num']
# Распознаем текст фрагмента
recognized_text = fragment['text']
current_line.append(recognized_text)
# Добавляем последнюю строку
if current_line:
result_lines.append(" ".join(current_line))
# Собираем результат в общий текст
final_text = "\n".join(result_lines)
if text_output_path:
with open(text_output_path, "w", encoding="utf-8") as text_file:
text_file.write(final_text)
return final_text
|