Overglitch commited on
Commit
2f2f4c9
·
verified ·
1 Parent(s): 78bf8ed

Update modules/preprocessing.py

Browse files
Files changed (1) hide show
  1. 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
- # 1) Cargar el PDF
102
- doc = DocumentFile.from_pdf(pdf_path)
103
-
104
- # 2) Crear un predictor (modelo OCR); docTR brinda modelos preentrenados
105
- predictor = ocr_predictor(pretrained=True)
106
-
107
- # 3) Aplicar el predictor al documento para obtener el layout
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: