Upload 5 files

README.md

@@ -26,10 +26,22 @@ print(generated_text)
 ```
 ## Usage for recognizing the page
 1. Download the following files from the repository: recognize_page.py, requirements_page.txt
-2. Use this comand
-```bash
-sudo apt install tesseract-ocr
-```
+2. Install tesseract-ocr
+    - Linux
+    ```bash
+    sudo apt install tesseract-ocr
+    ```
+    - MacOS
+    ```bash
+    brew install tesseract-ocr
+    ```
+    - Windows
+        1. Download tesseract exe from https://github.com/UB-Mannheim/tesseract/wiki.
+
+        2. Install this exe in C:\Program Files (x86)\Tesseract-OCR
+
+        3. Open virtual machine command prompt in windows or anaconda prompt.
+
 3. Install dependencies: 
 ```bash
 pip install -r requirements_page.txt
@@ -39,12 +51,41 @@ pip install -r requirements_page.txt
 from recognize_page import recognize_page
 
 page_path = "cleaned_pages/C2V10/page11.png"
-text = recognize_page(page_path)
+text = recognize_page(page_path, text_output_path="output/file.txt")
 
 print(f"Текст страницы:\n")
 print(text)
 ```
 
+## Usage for recognizing the book from pdf-file
+1. Download the following files from the repository: recognize_book.py, recognize_page.py, requirements_book.txt
+2. Install tesseract-ocr
+    - Linux
+    ```bash
+    sudo apt install tesseract-ocr
+    ```
+    - MacOS
+    ```bash
+    brew install tesseract-ocr
+    ```
+    - Windows
+        1. Download tesseract exe from https://github.com/UB-Mannheim/tesseract/wiki.
+
+        2. Install this exe in C:\Program Files (x86)\Tesseract-OCR
+
+        3. Open virtual machine command prompt in windows or anaconda prompt.
+
+3. Install dependencies: 
+```bash
+pip install -r requirements_book.txt
+```
+4. Use this code
+```python
+from recognize_book import recognize_book
+
+recognize_book(book_path="path/to/your/book.pdf", text_output_path="book_text.txt")
+```
+
 ## Metrics on test
 - CER (Char Error Rate) = 0.095
 - WER (Word Error Rate) = 0.298

recognize_book.py

@@ -0,0 +1,32 @@
+import tempfile
+import os
+from pdf2image import convert_from_path
+from recognize_page import recognize_page
+from tqdm import tqdm
+
+def recognize_book(book_path: str, text_output_path: str):
+    data = {}
+    
+    pages = convert_from_path(book_path)
+
+    for i in tqdm(range(len(pages)), desc="Procces pdf"):
+        with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as temp_file:
+            # Сохранение страницы в временный файл
+            pages[i].save(temp_file.name, 'PNG')
+            page_text = recognize_page(temp_file.name)
+            
+            data[i] = page_text
+            
+        os.remove(temp_file.name)
+    
+    book_text = ""
+    for i in data.keys():
+        book_text += f"\n\n=== Page {i+1} ===\n\n"
+        book_text += data[i] + f"\n"
+        
+    with open(text_output_path, "w", encoding="utf-8") as text_file:
+        text_file.write(book_text)
+
+
+if __name__ == "__main__":
+    recognize_book(book_path="bv000030992_0001.pdf", text_output_path="book_text.txt")

recognize_page.py

@@ -1,99 +1,112 @@
-import os
-from PIL import Image
 import pytesseract
-from tqdm import tqdm
-from pytesseract import Output
+import json
+import numpy as np
+from PIL import Image, ImageEnhance
 from transformers import VisionEncoderDecoderModel, TrOCRProcessor
 
-def recognize_row(row_file):
-    hf_model = VisionEncoderDecoderModel.from_pretrained("Serovvans/trocr-prereform-orthography")
-    image = Image.open(row_file)
-    processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-printed")
+
+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, output_dir="./", page_name=None):
-    """
-    Разбивает изображение страницы на строки, сортирует строки, распознаёт их и соединяет текст.
-
-    Параметры:
-        image_path (str): Путь к изображению страницы.
-        output_dir (str): Путь к папке для сохранения строк.
-        page_name (str): Имя страницы для сохранения строк (по умолчанию None).
 
-    Возвращает:
-        str: Итоговый распознанный текст страницы.
-    """
-    os.makedirs(output_dir, exist_ok=True)
-    image = Image.open(image_path)
-    data = pytesseract.image_to_data(image, config='--psm 3', output_type=Output.DICT)
-
-    lines = []
-    current_line = []
-    previous_y = None
-    y_threshold = 15  # Порог для объединения слов в одну строку
-
-    # Обход всех блоков текста
-    n_boxes = len(data['level'])
-    for i in range(n_boxes):
-        if data['level'][i] == 5 or data['level'][i] == 4:  # Уровень строки или слова
+# Основная функция распознавания текста с разделением на слова/фразы
+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]
-            text = data['text'][i].strip()
-
-            if not text:
-                continue
-
-            # Проверяем, начинается ли новая строка
-            if previous_y is None or abs(y - previous_y) > y_threshold:
-                if current_line:
-                    # Объединяем слова в строку и добавляем в список строк
-                    min_x = min([word['x'] for word in current_line])
-                    max_x = max([word['x'] + word['w'] for word in current_line])
-                    avg_y = sum([word['y'] for word in current_line]) / len(current_line)
-                    max_y = max([word['y'] + word['h'] for word in current_line])
-                    lines.append((min_x, avg_y, max_x - min_x, max_y - avg_y, current_line))
-                    current_line = []
-
-            # Добавляем текущее слово в текущую строку
-            current_line.append({'x': x, 'y': y, 'w': w, 'h': h, 'text': text})
-            previous_y = y
+            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:
-        min_x = min([word['x'] for word in current_line])
-        max_x = max([word['x'] + word['w'] for word in current_line])
-        avg_y = sum([word['y'] for word in current_line]) / len(current_line)
-        max_y = max([word['y'] + word['h'] for word in current_line])
-        lines.append((min_x, avg_y, max_x - min_x, max_y - avg_y, current_line))
-
-    # Сортировка строк по координате Y
-    lines.sort(key=lambda line: line[1])
-
-    # Сохранение строк как изображений и распознавание текста
-    recognized_text = []
-    i = 0
-    for line in tqdm(lines, desc="Processing page"):
-        x, y, w, h, words = line
-        min_x = x
-        max_x = x + w
-        min_y = max(0, y - 10)
-        max_y = y + h
-
-        # Обрезаем изображение по координатам строки
-        row_image = image.crop((min_x, min_y, max_x, max_y))
-        row_image_path = os.path.join(output_dir, f'{page_name}_row_{i}.png')
-        row_image.save(row_image_path)
-
-        # Распознаём текст строки
-        row_text = recognize_row(row_image_path)
-        
-        os.remove(row_image_path)
-        
-        recognized_text.append(row_text)
-        i += 1
-
-    # Соединяем распознанный текст строк
-    full_text = ' '.join(recognized_text)
-    return full_text
+        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

requirements_book.txt

@@ -0,0 +1,7 @@
+numpy==1.23.0
+pdf2image==1.17.0
+pytesseract==0.3.10
+Pillow==10.0.0
+transformers==4.33.2
+torch==2.0.1
+tqdm

requirements_page.txt

@@ -1,4 +1,6 @@
 pytesseract==0.3.10
 Pillow==10.0.0
 transformers==4.33.2
-torch==2.0.1
+torch==2.0.1
+numpy==1.23.0
+tqdm