app.py

from transformers import pipeline, TrOCRProcessor, VisionEncoderDecoderModel, T5ForConditionalGeneration, T5Tokenizer
from pdf2image import convert_from_path, convert_from_bytes
from IPython.display import clear_output
from PIL import Image
import cv2
import numpy as np
import torch
import gradio as gr

MIN_BOX_WIDTH = 8 # Минимальная ширина текстовой области (в пикселях)
MIN_BOX_HEIGHT = 15 # Минимальная высота текстовой области (в пикселях)
MAX_PART_WIDTH = 600 # Максимальная ширина части строки (в пикселях)
BOX_HEIGHT_TOLERANCE = 8 # Максимальная разница между высотами текстовых областей для добавлению в строку (в пикселях)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-printed")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-printed")
model.to(device)

summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=device)

model_translation = T5ForConditionalGeneration.from_pretrained('utrobinmv/t5_translate_en_ru_zh_small_1024')
model_translation.to(device)
tokenizer_translation = T5Tokenizer.from_pretrained('utrobinmv/t5_translate_en_ru_zh_small_1024')

def get_text_from_images(images):
  extracted_text = []
  image_number = 0
  for image in images:
    image_number += 1
    image_cv = np.array(image)
    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)

    gray = cv2.cvtColor(image_cv, cv2.COLOR_BGR2GRAY)
    thresh = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 11, 2)
    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    bounding_boxes = [cv2.boundingRect(contour) for contour in contours]

    def group_boxes_into_lines(boxes, tolerance=BOX_HEIGHT_TOLERANCE):
        sorted_boxes = sorted(boxes, key=lambda box: box[1])

        lines = []
        current_line = []

        for box in sorted_boxes:
            x, y, w, h = box

            if not current_line:
                current_line.append(box)
            else:
                last_box = current_line[-1]
                last_y = last_box[1]

                if abs(y - last_y) <= tolerance:
                    current_line.append(box)
                else:
                    lines.append(current_line)
                    current_line = [box]

        if current_line:
            lines.append(current_line)

        return lines

    lines = group_boxes_into_lines(bounding_boxes)

    line_number = 0
    for line in lines:
        line_number += 1

        x_coords = [box[0] for box in line]
        y_coords = [box[1] for box in line]
        widths = [box[2] for box in line]
        heights = [box[3] for box in line]

        x_min = min(x_coords)
        y_min = min(y_coords)
        x_max = max(x_coords[i] + widths[i] for i in range(len(line)))
        y_max = max(y_coords[i] + heights[i] for i in range(len(line)))

        line_image = image_cv[y_min:y_max, x_min:x_max]

        if line_image.size == 0 or line_image.shape[0] < MIN_BOX_HEIGHT or line_image.shape[1] < MIN_BOX_WIDTH:
            continue

        parts = []

        if line_image.shape[1] > MAX_PART_WIDTH:
          num_parts = (line_image.shape[1] // MAX_PART_WIDTH) + 1
          part_width = line_image.shape[1] // num_parts

          for i in range(num_parts):
            start_x = i * part_width
            end_x = (i + 1) * part_width if i < num_parts - 1 else line_image.shape[1]
            part = line_image[:, start_x:end_x]
            parts.append(part)
        else:
          parts.append(line_image)

        line_text = ""
        part_number = 0

        for part in parts:
          part_number += 1
          #clear_output()
          print(f"Images: {image_number}/{len(images)}")
          print(f"Lines: {line_number}/{len(lines)}")
          print(f"Parts: {part_number}/{len(parts)}")

          part_image_pil = Image.fromarray(cv2.cvtColor(part, cv2.COLOR_BGR2RGB))
          #display(part_image_pil)
          print("\n".join(extracted_text))

          pixel_values = processor(part_image_pil, return_tensors="pt").pixel_values
          pixel_values = pixel_values.to(device)
          generated_ids = model.generate(pixel_values)
          text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

          line_text += text

        extracted_text.append(line_text)

  final_text = "\n".join(extracted_text)
  return final_text

def summarize(text, max_length=300, min_length=150):
  result = summarizer(text, max_length=max_length, min_length=min_length, do_sample=False)
  return result[0]['summary_text']

def translate(text):
  prefix = 'translate to ru: '
  src_text = prefix + text

  input_ids = tokenizer_translation(src_text, return_tensors="pt")

  generated_tokens = model_translation.generate(**input_ids.to(device))

  result = tokenizer_translation.batch_decode(generated_tokens, skip_special_tokens=True)
  return result[0]

def launch(images, language):
  if images == None or not images:
    return "No input provided."
  raw_text = get_text_from_images(images)
  summary = summarize(raw_text)
  if language == "rus":
    return translate(summary)
  return summary

def pdf_to_image(pdf, index = 0):
  images = convert_from_bytes(pdf)
  if 0 <= index < len(images):
    return [images[index]]
  return []

def pdf_to_images(pdf):
  images = convert_from_bytes(pdf)
  return images

def process_pdf(pdf_file, process_mode, page_index, language): 
  if process_mode == "all":
    return launch(pdf_to_images(pdf_file), language)
  elif process_mode == "single":
    return launch(pdf_to_image(pdf_file, page_index), language)

def process_images(images, language):
  pil_images = []
  for image in images:
    pil_images.append(Image.open(image))
  launch(pil_images, language)

class PrintToTextbox:
    def __init__(self, textbox):
        self.textbox = textbox
        self.buffer = ""

    def write(self, text):
        self.buffer += text
        self.textbox.update(self.buffer)

    def flush(self):
        pass

def update_page_index_visibility(process_mode):
    if process_mode == "single":
        return gr.update(visible=True)
    else:
        return gr.update(visible=False)

with gr.Blocks() as demo:
    gr.Markdown("# PDF and Image Text Summarizer")
    gr.Markdown("Upload a PDF file or images to extract and summarize text.")
    gr.Markdown("Takes about 10 minutes per page.")

    language = gr.Radio(choices=["rus", "eng"], label="Output Language", value="rus")
    
    with gr.Tabs():
        with gr.TabItem("PDF"):
            pdf_file = gr.File(label="Upload PDF File", type="binary")
            process_mode = gr.Radio(choices=["single", "all"], label="Process Mode", value="single")
            page_index = gr.Number(label="Page Index", value=0, precision=0)
            pdf_output = gr.Textbox(label="Extracted Text")
            pdf_button = gr.Button("Extract Text from PDF")
        
        with gr.TabItem("Images"):
            images = gr.Files(label="Upload Images", file_types=["image"])
            image_output = gr.Textbox(label="Extracted Text")
            image_button = gr.Button("Extract Text from Images")
    
    pdf_button.click(process_pdf, inputs=[pdf_file, process_mode, page_index, language], outputs=pdf_output)
    image_button.click(process_images, inputs=[images, language], outputs=image_output)
    process_mode.change(update_page_index_visibility, inputs=process_mode, outputs=page_index)

demo.launch(debug=True)