mirror of
https://github.com/Ladebeze66/ragflow_preprocess.git
synced 2026-02-04 05:50:26 +01:00
138 lines
5.0 KiB
Python
138 lines
5.0 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
Traitement des PDF avec le gestionnaire de flux de travail
|
|
"""
|
|
|
|
import os
|
|
import io
|
|
import time
|
|
from typing import Dict, Any, Optional, List, Tuple, Union
|
|
from PIL import Image
|
|
|
|
from utils.workflow_manager import WorkflowManager
|
|
|
|
class PDFProcessor:
|
|
"""
|
|
Processeur de PDF intégré avec le gestionnaire de flux de travail
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""
|
|
Initialise le processeur de PDF
|
|
"""
|
|
# Initialiser le gestionnaire de flux
|
|
self.workflow_manager = WorkflowManager()
|
|
|
|
# Dossier pour sauvegarder les résultats
|
|
self.output_dir = os.path.join("data", "outputs")
|
|
os.makedirs(self.output_dir, exist_ok=True)
|
|
|
|
# Timestamp pour cette session
|
|
self.timestamp = time.strftime("%Y%m%d-%H%M%S")
|
|
|
|
print(f"Processeur PDF initialisé")
|
|
|
|
def process_image_selection(self, image_data: bytes, selection_type: str,
|
|
context: str, page_number: int) -> Dict[str, Any]:
|
|
"""
|
|
Traite une sélection d'image dans un PDF
|
|
|
|
Args:
|
|
image_data (bytes): Données de l'image sélectionnée
|
|
selection_type (str): Type de sélection (schéma, tableau, etc.)
|
|
context (str): Contexte textuel associé à l'image
|
|
page_number (int): Numéro de la page où se trouve l'image
|
|
|
|
Returns:
|
|
Dict: Résultats du traitement
|
|
"""
|
|
print(f"Traitement d'une image de type '{selection_type}' (page {page_number})")
|
|
|
|
# Traiter l'image avec le gestionnaire de flux
|
|
results = self.workflow_manager.process_image_with_context(
|
|
image_data, selection_type, context
|
|
)
|
|
|
|
# Ajouter des informations supplémentaires aux résultats
|
|
results["page_number"] = page_number
|
|
results["timestamp"] = time.strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
# Sauvegarder les résultats dans un fichier texte
|
|
self._save_results(results, selection_type, page_number)
|
|
|
|
return results
|
|
|
|
def _save_results(self, results: Dict[str, Any], selection_type: str, page_number: int) -> str:
|
|
"""
|
|
Sauvegarde les résultats du traitement dans un fichier
|
|
|
|
Args:
|
|
results (Dict): Résultats du traitement
|
|
selection_type (str): Type de sélection
|
|
page_number (int): Numéro de page
|
|
|
|
Returns:
|
|
str: Chemin du fichier de résultats
|
|
"""
|
|
# Créer un nom de fichier unique
|
|
filename = f"{self.timestamp}_page{page_number}_{selection_type}.txt"
|
|
output_path = os.path.join(self.output_dir, filename)
|
|
|
|
# Sauvegarder l'image si disponible
|
|
if "original" in results and "image" in results["original"]:
|
|
try:
|
|
image_data = results["original"]["image"]
|
|
img = Image.open(io.BytesIO(image_data))
|
|
|
|
image_filename = f"{self.timestamp}_page{page_number}_{selection_type}.png"
|
|
image_path = os.path.join(self.output_dir, image_filename)
|
|
|
|
img.save(image_path)
|
|
print(f"Image sauvegardée dans: {image_path}")
|
|
except Exception as e:
|
|
print(f"Erreur lors de la sauvegarde de l'image: {str(e)}")
|
|
|
|
# Écrire les résultats dans un fichier texte
|
|
with open(output_path, "w", encoding="utf-8") as f:
|
|
# En-tête
|
|
f.write(f"# Analyse de {selection_type} - Page {page_number}\n")
|
|
f.write(f"Date et heure: {results['timestamp']}\n\n")
|
|
|
|
# Contexte original
|
|
if "original" in results and "context" in results["original"]:
|
|
f.write("## Contexte original (français)\n\n")
|
|
f.write(f"{results['original']['context']}\n\n")
|
|
|
|
# Contexte traduit
|
|
if "context_en" in results:
|
|
f.write("## Contexte traduit (anglais)\n\n")
|
|
f.write(f"{results['context_en']}\n\n")
|
|
|
|
# Analyse de vision
|
|
if "vision_analysis" in results:
|
|
f.write("## Analyse de l'image (anglais)\n\n")
|
|
f.write(f"{results['vision_analysis']}\n\n")
|
|
|
|
# Analyse traduite
|
|
if "analysis_fr" in results:
|
|
f.write("## Analyse traduite (français)\n\n")
|
|
f.write(f"{results['analysis_fr']}\n\n")
|
|
|
|
# Erreurs éventuelles
|
|
if "error" in results:
|
|
f.write("## ERREUR\n\n")
|
|
f.write(f"{results['error']}\n\n")
|
|
|
|
print(f"Résultats sauvegardés dans: {output_path}")
|
|
return output_path
|
|
|
|
def get_workflow_log_path(self) -> str:
|
|
"""
|
|
Renvoie le chemin du fichier de journal du flux de travail
|
|
|
|
Returns:
|
|
str: Chemin du fichier de journal
|
|
"""
|
|
return self.workflow_manager.logger.get_workflow_path() |