358 lines
15 KiB
Python
358 lines
15 KiB
Python
# backend/utils.py
|
|
import os
|
|
import time
|
|
import pytesseract
|
|
from PIL import Image
|
|
from pdf2image import convert_from_path
|
|
import shutil
|
|
import asyncio
|
|
import uuid
|
|
import subprocess
|
|
import base64
|
|
import re
|
|
from models import get_prompt
|
|
|
|
from reportlab.lib.pagesizes import A4
|
|
from reportlab.platypus import SimpleDocTemplate, Paragraph, Spacer, PageBreak
|
|
from reportlab.lib.styles import getSampleStyleSheet, ParagraphStyle
|
|
from reportlab.lib.enums import TA_JUSTIFY, TA_CENTER
|
|
|
|
# Se agrega BeautifulSoup para procesar HTML
|
|
from bs4 import BeautifulSoup
|
|
|
|
UPLOAD_FOLDER = "uploads"
|
|
OUTPUT_FOLDER = "outputs"
|
|
|
|
def encode_image(file_path):
|
|
with open(file_path, "rb") as f:
|
|
return base64.b64encode(f.read()).decode("utf-8")
|
|
|
|
def run_tesseract(file_path):
|
|
"""
|
|
If the file is a PDF, perform OCR page by page and add a "Page X:" header;
|
|
otherwise, perform OCR normally.
|
|
"""
|
|
extracted_text = ""
|
|
if file_path.lower().endswith(".pdf"):
|
|
try:
|
|
pages = convert_from_path(file_path)
|
|
for i, page in enumerate(pages, start=1):
|
|
page_text = pytesseract.image_to_string(page, lang='eng')
|
|
extracted_text += f"Page {i}:\n{page_text}\n\n"
|
|
except Exception as e:
|
|
extracted_text = f"❌ Error processing PDF: {str(e)}"
|
|
else:
|
|
image = Image.open(file_path)
|
|
extracted_text = pytesseract.image_to_string(image, lang='eng')
|
|
return extracted_text
|
|
|
|
def call_api_correction(api, model, text, prompt_key="ocr_correction"):
|
|
prompt = get_prompt(prompt_key) + text
|
|
try:
|
|
from google import genai
|
|
except ImportError:
|
|
raise ImportError("Please install 'google-genai' to use Gemini.")
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
|
|
response = client.models.generate_content(model=model, contents=[prompt])
|
|
return response.text
|
|
|
|
def call_api_ocr(api, model, file_path, prompt_key="ocr"):
|
|
try:
|
|
from google import genai
|
|
except ImportError:
|
|
raise ImportError("Please install 'google-genai' to use Gemini.")
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
|
|
file_ref = client.files.upload(file=file_path)
|
|
prompt = get_prompt(prompt_key)
|
|
response = client.models.generate_content(model=model, contents=[file_ref, prompt])
|
|
return response.text
|
|
|
|
def call_api_translation(api, model, text, target_language, prompt_key="translation"):
|
|
prompt_template = get_prompt(prompt_key)
|
|
prompt = prompt_template.format(target_language=target_language) + text
|
|
try:
|
|
from google import genai
|
|
except ImportError:
|
|
raise ImportError("Please install 'google-genai' to use Gemini.")
|
|
try:
|
|
loop = asyncio.get_event_loop()
|
|
except RuntimeError:
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
|
|
response = client.models.generate_content(model=model, contents=[prompt])
|
|
return response.text
|
|
|
|
def embed_ocr_in_pdf(input_pdf, output_pdf):
|
|
try:
|
|
subprocess.run(["ocrmypdf", input_pdf, output_pdf], check=True)
|
|
return True
|
|
except Exception as e:
|
|
return False
|
|
|
|
def ocr_file_by_pages(file_path, api, model, prompt_key, update_progress, is_cancelled):
|
|
final_text = ""
|
|
if file_path.lower().endswith(".pdf"):
|
|
try:
|
|
pages = convert_from_path(file_path)
|
|
except Exception as e:
|
|
return f"❌ Error processing PDF: {str(e)}"
|
|
total = len(pages)
|
|
for i, page in enumerate(pages, start=1):
|
|
if is_cancelled():
|
|
update_progress(0, "⏹️ Cancelled")
|
|
return "Process cancelled."
|
|
temp_filename = os.path.join(OUTPUT_FOLDER, f"temp_page_{uuid.uuid4().hex}.png")
|
|
page.save(temp_filename, "PNG")
|
|
page_text = call_api_ocr(api, model, temp_filename, prompt_key)
|
|
final_text += f"Page {i}:\n{page_text}\n\n"
|
|
os.remove(temp_filename)
|
|
progress = int((i / total) * 100)
|
|
update_progress(progress, f"📄 Processed page {i} of {total}.")
|
|
time.sleep(1)
|
|
return final_text
|
|
else:
|
|
return call_api_ocr(api, model, file_path, prompt_key)
|
|
|
|
def translate_file_by_pages(file_path, api, model, target_language, prompt_key, update_progress, is_cancelled):
|
|
final_translation = ""
|
|
if file_path.lower().endswith(".pdf"):
|
|
try:
|
|
pages = convert_from_path(file_path)
|
|
except Exception as e:
|
|
return f"❌ Error processing PDF: {str(e)}"
|
|
total = len(pages)
|
|
for i, page in enumerate(pages, start=1):
|
|
if is_cancelled():
|
|
update_progress(0, "⏹️ Cancelled")
|
|
return "Process cancelled."
|
|
temp_filename = os.path.join(OUTPUT_FOLDER, f"temp_page_{uuid.uuid4().hex}.png")
|
|
page.save(temp_filename, "PNG")
|
|
page_text = pytesseract.image_to_string(page, lang='eng')
|
|
translated_page = call_api_translation(api, model, page_text, target_language, prompt_key)
|
|
final_translation += f"Page {i}:\n{translated_page}\n\n"
|
|
os.remove(temp_filename)
|
|
progress = int((i / total) * 100)
|
|
update_progress(progress, f"📄 Processed page {i} of {total}.")
|
|
time.sleep(1)
|
|
return final_translation
|
|
elif file_path.lower().endswith(".txt"):
|
|
with open(file_path, "r", encoding="utf-8") as f:
|
|
text = f.read()
|
|
translated = call_api_translation(api, model, text, target_language, prompt_key)
|
|
update_progress(100, "🎉 Process completed")
|
|
return f"Page 1:\n{translated}"
|
|
else:
|
|
return "Unsupported file type for translation."
|
|
|
|
def process_file(file_path, api, model, mode, prompt_key, update_progress, is_cancelled):
|
|
if is_cancelled():
|
|
update_progress(0, "⏹️ Cancelled")
|
|
return "Process cancelled."
|
|
update_progress(25, "📤 File uploaded.")
|
|
base_name = os.path.splitext(os.path.basename(file_path))[0]
|
|
|
|
if mode == "OCR":
|
|
# Process using tesseract with page-structure if PDF.
|
|
processed_text = run_tesseract(file_path)
|
|
if is_cancelled():
|
|
update_progress(25, "⏹️ Cancelled")
|
|
return "Process cancelled."
|
|
update_progress(50, "✅ Tesseract OCR completed.")
|
|
# Generate PDF copy as before.
|
|
pdf_output = os.path.join(OUTPUT_FOLDER, base_name + "_ocr.pdf")
|
|
if file_path.lower().endswith(".pdf"):
|
|
if embed_ocr_in_pdf(file_path, pdf_output):
|
|
update_progress(95, "📄 OCR embedded into PDF.")
|
|
else:
|
|
shutil.copy(file_path, pdf_output)
|
|
update_progress(95, "⚠️ Failed to embed OCR; original PDF copied.")
|
|
else:
|
|
shutil.copy(file_path, pdf_output)
|
|
elif mode == "OCR + AI":
|
|
processed_text = run_tesseract(file_path)
|
|
if is_cancelled():
|
|
update_progress(25, "⏹️ Cancelled")
|
|
return "Process cancelled."
|
|
update_progress(50, "✅ Tesseract OCR completed.")
|
|
processed_text = call_api_correction(api, model, processed_text, prompt_key)
|
|
pdf_output = os.path.join(OUTPUT_FOLDER, base_name + "_ocr.pdf")
|
|
if file_path.lower().endswith(".pdf"):
|
|
if embed_ocr_in_pdf(file_path, pdf_output):
|
|
update_progress(95, "📄 OCR embedded into PDF.")
|
|
else:
|
|
shutil.copy(file_path, pdf_output)
|
|
update_progress(95, "⚠️ Failed to embed OCR; original PDF copied.")
|
|
else:
|
|
shutil.copy(file_path, pdf_output)
|
|
elif mode == "AI":
|
|
update_progress(25, "📂 File ready for full AI processing.")
|
|
if file_path.lower().endswith(".pdf"):
|
|
processed_text = ocr_file_by_pages(file_path, api, model, prompt_key, update_progress, is_cancelled)
|
|
else:
|
|
processed_text = call_api_ocr(api, model, file_path, prompt_key)
|
|
# In AI mode, do not generate a new PDF.
|
|
else:
|
|
processed_text = "Unrecognized processing mode."
|
|
update_progress(25, "❌ Error: Unrecognized mode.")
|
|
|
|
if is_cancelled():
|
|
update_progress(75, "⏹️ Cancelled")
|
|
return "Process cancelled."
|
|
|
|
update_progress(75, "🤖 API processing completed.")
|
|
|
|
# Write the output TXT file (it is already structured by page in OCR and OCR+AI modes)
|
|
txt_file = os.path.join(OUTPUT_FOLDER, base_name + ".txt")
|
|
with open(txt_file, "w", encoding="utf-8") as f:
|
|
f.write(processed_text)
|
|
update_progress(100, "🎉 Process completed")
|
|
return processed_text
|
|
|
|
def process_text(text):
|
|
lines = text.splitlines()
|
|
processed_lines = []
|
|
buffer = ""
|
|
for line in lines:
|
|
stripped = line.strip()
|
|
if not stripped:
|
|
if buffer:
|
|
processed_lines.append(buffer)
|
|
buffer = ""
|
|
else:
|
|
processed_lines.append("")
|
|
continue
|
|
if buffer:
|
|
if buffer.endswith('.'):
|
|
processed_lines.append(buffer)
|
|
buffer = stripped
|
|
else:
|
|
buffer += " " + stripped
|
|
else:
|
|
buffer = stripped
|
|
if buffer:
|
|
processed_lines.append(buffer)
|
|
return "\n".join(processed_lines)
|
|
|
|
def convert_txt_to_pdf(txt_file_path):
|
|
"""
|
|
Se ha modificado para que, si el contenido del TXT está formateado en HTML,
|
|
se interpreten los tags que indiquen títulos (<h1>, <h2>, etc.), párrafos (<p>)
|
|
y saltos de página (<pagebreak> o div con clase "page-break").
|
|
|
|
En caso de que el contenido sea texto plano y contenga patrones en el formato
|
|
[Page X] (entre corchetes), se usará ese separador para dividir las páginas.
|
|
El encabezado (sin los corchetes) se incluirá en la parte superior de cada página.
|
|
"""
|
|
with open(txt_file_path, "r", encoding="utf-8") as f:
|
|
content = f.read()
|
|
|
|
base_name = os.path.splitext(os.path.basename(txt_file_path))[0]
|
|
output_pdf = os.path.join(OUTPUT_FOLDER, base_name + "_txt.pdf")
|
|
doc = SimpleDocTemplate(output_pdf, pagesize=A4,
|
|
rightMargin=40, leftMargin=40,
|
|
topMargin=40, bottomMargin=40)
|
|
styles = getSampleStyleSheet()
|
|
# Estilos para encabezados
|
|
header_styles = {
|
|
"h1": ParagraphStyle('Heading1', parent=styles['Heading1'], alignment=TA_CENTER),
|
|
"h2": ParagraphStyle('Heading2', parent=styles['Heading2'], alignment=TA_CENTER),
|
|
"h3": ParagraphStyle('Heading3', parent=styles['Heading3'], alignment=TA_CENTER),
|
|
"h4": ParagraphStyle('Heading4', parent=styles['Heading4'], alignment=TA_CENTER),
|
|
"h5": ParagraphStyle('Heading5', parent=styles['Heading5'], alignment=TA_CENTER),
|
|
"h6": ParagraphStyle('Heading6', parent=styles['Heading6'], alignment=TA_CENTER),
|
|
}
|
|
normal_style = ParagraphStyle(
|
|
'Normal',
|
|
parent=styles['Normal'],
|
|
alignment=TA_JUSTIFY,
|
|
leading=15,
|
|
leftIndent=20 # Sangría al inicio de cada párrafo
|
|
)
|
|
|
|
flowables = []
|
|
# Detectamos si el contenido es HTML (buscando etiquetas comunes)
|
|
is_html = any(tag in content.lower() for tag in ["<html", "<p", "<h1", "<h2", "<h3"])
|
|
|
|
if is_html:
|
|
soup = BeautifulSoup(content, "html.parser")
|
|
body = soup.body if soup.body else soup
|
|
for element in body.children:
|
|
if element.name is None:
|
|
text = element.strip()
|
|
if text:
|
|
para = Paragraph(text, normal_style)
|
|
flowables.append(para)
|
|
flowables.append(Spacer(1, 12))
|
|
elif element.name.lower() in ["h1", "h2", "h3", "h4", "h5", "h6"]:
|
|
tag = element.name.lower()
|
|
style = header_styles.get(tag, styles['Heading1'])
|
|
para = Paragraph(element.decode_contents(), style)
|
|
flowables.append(para)
|
|
flowables.append(Spacer(1, 12))
|
|
elif element.name.lower() == "p":
|
|
para = Paragraph(element.decode_contents(), normal_style)
|
|
flowables.append(para)
|
|
flowables.append(Spacer(1, 12))
|
|
elif element.name.lower() == "pagebreak" or (element.name.lower() == "div" and "page-break" in element.get("class", [])):
|
|
flowables.append(PageBreak())
|
|
else:
|
|
text = element.get_text().strip()
|
|
if text:
|
|
para = Paragraph(text, normal_style)
|
|
flowables.append(para)
|
|
flowables.append(Spacer(1, 12))
|
|
else:
|
|
# Si se detecta el patrón [Page X] en el contenido, se usa para dividir las páginas.
|
|
if re.search(r'\[Page\s+\d+\]', content, re.IGNORECASE):
|
|
# Dividir incluyendo el separador (usamos grupo de captura)
|
|
parts = re.split(r'(\[Page\s+\d+\])', content, flags=re.IGNORECASE)
|
|
current_header = ""
|
|
for part in parts:
|
|
part = part.strip()
|
|
if not part:
|
|
continue
|
|
# Si es un encabezado [Page X]
|
|
if re.match(r'\[Page\s+\d+\]', part, re.IGNORECASE):
|
|
# Extraemos el número o texto sin los corchetes para usarlo como encabezado
|
|
current_header = part.strip("[]")
|
|
header_para = Paragraph(current_header, styles['Heading1'])
|
|
flowables.append(header_para)
|
|
flowables.append(Spacer(1, 12))
|
|
else:
|
|
# Es el contenido de la página
|
|
processed = process_text(part)
|
|
para = Paragraph(processed.replace("\n", "<br/>"), normal_style)
|
|
flowables.append(para)
|
|
flowables.append(Spacer(1, 12))
|
|
flowables.append(PageBreak())
|
|
else:
|
|
# Si no se detectan separadores, se usa el método anterior basado en "Page X:" sin corchetes.
|
|
pages = re.split(r'(?i)Page\s+\d+:\s*', content)
|
|
if pages and pages[0].strip() == "":
|
|
pages = pages[1:]
|
|
if len(pages) <= 1:
|
|
pages = content.split("\n\n")
|
|
processed_pages = [process_text(page) for page in pages if page.strip() != ""]
|
|
for i, page_text in enumerate(processed_pages, start=1):
|
|
header = Paragraph(f"PAGE {i}", styles['Heading1'])
|
|
flowables.append(header)
|
|
flowables.append(Spacer(1, 12))
|
|
para = Paragraph(page_text.replace("\n", "<br/>"), normal_style)
|
|
flowables.append(para)
|
|
if i < len(processed_pages):
|
|
flowables.append(Spacer(1, 24))
|
|
flowables.append(PageBreak())
|
|
|
|
doc.build(flowables)
|
|
return output_pdf
|