Spaces:
Sleeping
Sleeping
Update modules/preprocessing.py
Browse files- modules/preprocessing.py +46 -21
modules/preprocessing.py
CHANGED
@@ -97,28 +97,53 @@ class PDFProcessor:
|
|
97 |
self.ocr_model = ocr_predictor(pretrained=True)
|
98 |
self.max_pages = max_pages
|
99 |
|
100 |
-
def pdf_to_text(pdf_path):
|
101 |
-
|
102 |
-
|
103 |
-
|
104 |
-
|
105 |
-
|
106 |
-
|
107 |
-
|
108 |
-
ocr_result = predictor(doc)
|
109 |
-
|
110 |
-
# Ahora sí, las páginas tienen .blocks, .lines, etc.
|
111 |
-
pages = ocr_result.pages
|
112 |
-
|
113 |
-
# 4) Extraer el texto de cada bloque
|
114 |
-
text_pages = []
|
115 |
-
for page in pages:
|
116 |
-
for block in page.blocks:
|
117 |
-
text_pages.append(block.text)
|
118 |
-
|
119 |
-
# 5) Unir o procesar a conveniencia
|
120 |
-
return "\n".join(text_pages)
|
121 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
122 |
|
123 |
|
124 |
class FileHandler:
|
|
|
97 |
self.ocr_model = ocr_predictor(pretrained=True)
|
98 |
self.max_pages = max_pages
|
99 |
|
100 |
+
def pdf_to_text(self, pdf_path):
|
101 |
+
"""
|
102 |
+
Convierte un archivo PDF a texto usando OCR.
|
103 |
+
Si el archivo no está en disco, lo guarda temporalmente.
|
104 |
+
"""
|
105 |
+
# Asegurarse de que el archivo está disponible temporalmente
|
106 |
+
temp_dir = Path("temp")
|
107 |
+
temp_dir.mkdir(exist_ok=True)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
108 |
|
109 |
+
# Manejar el archivo temporalmente
|
110 |
+
temp_file_path = temp_dir / Path(pdf_path).name
|
111 |
+
if not temp_file_path.exists():
|
112 |
+
shutil.copy(pdf_path, temp_file_path)
|
113 |
+
|
114 |
+
try:
|
115 |
+
# 1) Cargar el PDF
|
116 |
+
doc = DocumentFile.from_pdf(temp_file_path)
|
117 |
+
|
118 |
+
# 2) Limitar el número de páginas si es necesario
|
119 |
+
if len(doc.pages) > self.max_pages:
|
120 |
+
doc.pages = doc.pages[:self.max_pages]
|
121 |
+
|
122 |
+
# 3) Aplicar el modelo OCR al documento
|
123 |
+
ocr_result = self.ocr_model(doc)
|
124 |
+
|
125 |
+
# 4) Extraer texto de cada bloque
|
126 |
+
text_pages = []
|
127 |
+
for page in ocr_result.pages:
|
128 |
+
for block in page.blocks:
|
129 |
+
text_pages.append(block.text)
|
130 |
+
|
131 |
+
# 5) Unir todo el texto y devolverlo
|
132 |
+
return "\n".join(text_pages)
|
133 |
+
|
134 |
+
finally:
|
135 |
+
# Limpiar el archivo temporal después del procesamiento
|
136 |
+
if temp_file_path.exists():
|
137 |
+
temp_file_path.unlink()
|
138 |
+
|
139 |
+
@staticmethod
|
140 |
+
def clear_temp_directory():
|
141 |
+
"""Limpia todos los archivos en el directorio temporal."""
|
142 |
+
temp_dir = Path("temp")
|
143 |
+
if temp_dir.exists():
|
144 |
+
for file in temp_dir.iterdir():
|
145 |
+
if file.is_file():
|
146 |
+
file.unlink()
|
147 |
|
148 |
|
149 |
class FileHandler:
|