mirror of
https://github.com/Ladebeze66/ragflow_preprocess.git
synced 2026-02-04 05:50:26 +01:00
394 lines
15 KiB
Python
394 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding: utf-8 -*-
|
|
|
|
"""
|
|
Interface de visualisation et de sélection dans les documents PDF
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
from PyQt6.QtWidgets import (QMainWindow, QFileDialog, QVBoxLayout, QHBoxLayout,
|
|
QLabel, QPushButton, QComboBox, QWidget, QScrollArea,
|
|
QSpinBox, QFrame, QSplitter, QGroupBox, QStatusBar)
|
|
from PyQt6.QtCore import Qt, QRectF, QPoint, pyqtSignal
|
|
from PyQt6.QtGui import QPixmap, QPainter, QPen, QColor, QImage
|
|
|
|
from .llm_config_panel import LLMConfigPanel
|
|
import fitz # PyMuPDF
|
|
|
|
class PDFViewer(QMainWindow):
|
|
"""Interface principale pour la visualisation et l'annotation de PDFs"""
|
|
|
|
def __init__(self):
|
|
super().__init__()
|
|
|
|
self.pdf_document = None
|
|
self.current_page = 0
|
|
self.total_pages = 0
|
|
self.zoom_factor = 1.0
|
|
self.selection_rect = None
|
|
self.selection_start = None
|
|
self.selection_active = False
|
|
self.selected_regions = [] # Liste des zones sélectionnées
|
|
|
|
self.init_ui()
|
|
|
|
def init_ui(self):
|
|
"""Initialise l'interface utilisateur"""
|
|
self.setWindowTitle("Prétraitement PDF pour Ragflow")
|
|
self.setGeometry(100, 100, 1200, 800)
|
|
|
|
# Barre d'état
|
|
self.status_bar = QStatusBar()
|
|
self.setStatusBar(self.status_bar)
|
|
self.status_bar.showMessage("Prêt")
|
|
|
|
# Widget principal et layout
|
|
main_widget = QWidget()
|
|
self.setCentralWidget(main_widget)
|
|
main_layout = QHBoxLayout(main_widget)
|
|
|
|
# Splitter pour séparer la visualisation PDF et le panneau de configuration
|
|
splitter = QSplitter(Qt.Orientation.Horizontal)
|
|
main_layout.addWidget(splitter)
|
|
|
|
# =================== Partie gauche: Visualisation PDF ====================
|
|
pdf_panel = QWidget()
|
|
pdf_layout = QVBoxLayout(pdf_panel)
|
|
|
|
# Barre d'outils pour le PDF
|
|
toolbar = QWidget()
|
|
toolbar_layout = QHBoxLayout(toolbar)
|
|
toolbar_layout.setContentsMargins(0, 0, 0, 0)
|
|
|
|
# Bouton pour charger un PDF
|
|
self.load_btn = QPushButton("Charger PDF")
|
|
self.load_btn.clicked.connect(self.load_pdf)
|
|
toolbar_layout.addWidget(self.load_btn)
|
|
|
|
# Navigation
|
|
self.prev_btn = QPushButton("Page précédente")
|
|
self.prev_btn.clicked.connect(self.prev_page)
|
|
self.prev_btn.setEnabled(False)
|
|
toolbar_layout.addWidget(self.prev_btn)
|
|
|
|
self.page_spin = QSpinBox()
|
|
self.page_spin.setMinimum(1)
|
|
self.page_spin.setMaximum(1)
|
|
self.page_spin.valueChanged.connect(self.go_to_page)
|
|
toolbar_layout.addWidget(self.page_spin)
|
|
|
|
self.page_count_label = QLabel(" / 1")
|
|
toolbar_layout.addWidget(self.page_count_label)
|
|
|
|
self.next_btn = QPushButton("Page suivante")
|
|
self.next_btn.clicked.connect(self.next_page)
|
|
self.next_btn.setEnabled(False)
|
|
toolbar_layout.addWidget(self.next_btn)
|
|
|
|
# Zoom
|
|
self.zoom_in_btn = QPushButton("Zoom +")
|
|
self.zoom_in_btn.clicked.connect(self.zoom_in)
|
|
toolbar_layout.addWidget(self.zoom_in_btn)
|
|
|
|
self.zoom_out_btn = QPushButton("Zoom -")
|
|
self.zoom_out_btn.clicked.connect(self.zoom_out)
|
|
toolbar_layout.addWidget(self.zoom_out_btn)
|
|
|
|
# Bouton pour effacer la sélection
|
|
self.clear_sel_btn = QPushButton("Effacer sélection")
|
|
self.clear_sel_btn.clicked.connect(self.clear_selection)
|
|
toolbar_layout.addWidget(self.clear_sel_btn)
|
|
|
|
pdf_layout.addWidget(toolbar)
|
|
|
|
# Zone de visualisation du PDF
|
|
self.scroll_area = QScrollArea()
|
|
self.scroll_area.setWidgetResizable(True)
|
|
self.scroll_area.setAlignment(Qt.AlignmentFlag.AlignCenter)
|
|
|
|
self.pdf_label = PDFLabel(self)
|
|
self.scroll_area.setWidget(self.pdf_label)
|
|
|
|
pdf_layout.addWidget(self.scroll_area)
|
|
|
|
# =================== Partie droite: Panel de configuration ====================
|
|
self.config_panel = LLMConfigPanel(self)
|
|
|
|
# Ajout des deux panneaux au splitter
|
|
splitter.addWidget(pdf_panel)
|
|
splitter.addWidget(self.config_panel)
|
|
|
|
# Définir les proportions de départ du splitter
|
|
splitter.setSizes([700, 500])
|
|
|
|
def load_pdf(self):
|
|
"""Ouvre un dialogue pour charger un fichier PDF"""
|
|
file_path, _ = QFileDialog.getOpenFileName(
|
|
self, "Ouvrir un fichier PDF", "", "Fichiers PDF (*.pdf)"
|
|
)
|
|
|
|
if file_path:
|
|
try:
|
|
self.pdf_document = fitz.open(file_path)
|
|
self.total_pages = len(self.pdf_document)
|
|
self.current_page = 0
|
|
|
|
# Mise à jour de l'interface
|
|
self.page_spin.setMaximum(self.total_pages)
|
|
self.page_count_label.setText(f" / {self.total_pages}")
|
|
self.page_spin.setValue(1) # Cette action déclenche go_to_page()
|
|
|
|
self.prev_btn.setEnabled(True)
|
|
self.next_btn.setEnabled(True)
|
|
|
|
self.render_current_page()
|
|
|
|
# Mise à jour du titre avec le nom du fichier
|
|
file_name = os.path.basename(file_path)
|
|
self.setWindowTitle(f"Prétraitement PDF - {file_name}")
|
|
self.status_bar.showMessage(f"PDF chargé: {file_name}, {self.total_pages} pages")
|
|
|
|
except Exception as e:
|
|
self.status_bar.showMessage(f"Erreur lors du chargement du PDF: {str(e)}")
|
|
|
|
def render_current_page(self):
|
|
"""Affiche la page courante du PDF"""
|
|
if not self.pdf_document:
|
|
return
|
|
|
|
# Récupérer la page actuelle
|
|
page = self.pdf_document[self.current_page]
|
|
|
|
# Facteur de zoom et de qualité pour le rendu
|
|
zoom_matrix = fitz.Matrix(2 * self.zoom_factor, 2 * self.zoom_factor)
|
|
|
|
# Rendu de la page en un pixmap
|
|
try:
|
|
# Essayer avec la nouvelle API (PyMuPDF récent)
|
|
pixmap = page.get_pixmap(matrix=zoom_matrix, alpha=False)
|
|
except AttributeError:
|
|
# Fallback pour les anciennes versions
|
|
pixmap = page.render_pixmap(matrix=zoom_matrix, alpha=False)
|
|
|
|
# Conversion en QImage puis QPixmap
|
|
img = QImage(pixmap.samples, pixmap.width, pixmap.height,
|
|
pixmap.stride, QImage.Format.Format_RGB888)
|
|
|
|
pixmap_qt = QPixmap.fromImage(img)
|
|
|
|
# Mise à jour de l'affichage
|
|
self.pdf_label.setPixmap(pixmap_qt)
|
|
self.pdf_label.adjustSize()
|
|
|
|
# Mise à jour de l'état
|
|
self.status_bar.showMessage(f"Page {self.current_page + 1}/{self.total_pages}")
|
|
|
|
def next_page(self):
|
|
"""Passe à la page suivante"""
|
|
if self.pdf_document and self.current_page < self.total_pages - 1:
|
|
self.current_page += 1
|
|
self.page_spin.setValue(self.current_page + 1)
|
|
|
|
def prev_page(self):
|
|
"""Passe à la page précédente"""
|
|
if self.pdf_document and self.current_page > 0:
|
|
self.current_page -= 1
|
|
self.page_spin.setValue(self.current_page + 1)
|
|
|
|
def go_to_page(self, page_num):
|
|
"""Va à une page spécifique"""
|
|
if self.pdf_document and 1 <= page_num <= self.total_pages:
|
|
self.current_page = page_num - 1
|
|
self.render_current_page()
|
|
|
|
def zoom_in(self):
|
|
"""Augmente le facteur de zoom"""
|
|
self.zoom_factor *= 1.25
|
|
self.render_current_page()
|
|
|
|
def zoom_out(self):
|
|
"""Diminue le facteur de zoom"""
|
|
self.zoom_factor *= 0.8
|
|
self.render_current_page()
|
|
|
|
def clear_selection(self):
|
|
"""Efface la sélection actuelle"""
|
|
self.selection_rect = None
|
|
self.selection_active = False
|
|
self.selection_start = None
|
|
self.pdf_label.update()
|
|
|
|
def add_selection(self, rect, page_num=None):
|
|
"""
|
|
Ajoute une sélection à la liste des régions sélectionnées
|
|
|
|
Args:
|
|
rect (QRectF): Rectangle de sélection
|
|
page_num (int, optional): Numéro de page. Si None, utilise la page courante.
|
|
"""
|
|
page = self.current_page if page_num is None else page_num
|
|
|
|
# DEBUG: Afficher les coordonnées brutes et le facteur de zoom
|
|
print(f"DEBUG: Coordonnées brutes de sélection: x={rect.x()}, y={rect.y()}, w={rect.width()}, h={rect.height()}")
|
|
print(f"DEBUG: Facteur de zoom: {self.zoom_factor}")
|
|
|
|
# Calculer le décalage du viewport
|
|
viewport_pos = self.scroll_area.horizontalScrollBar().value(), self.scroll_area.verticalScrollBar().value()
|
|
print(f"DEBUG: Décalage du viewport: x={viewport_pos[0]}, y={viewport_pos[1]}")
|
|
|
|
# Obtenir les dimensions actuelles de l'image affichée
|
|
if self.pdf_label.pixmap():
|
|
img_width = self.pdf_label.pixmap().width()
|
|
img_height = self.pdf_label.pixmap().height()
|
|
print(f"DEBUG: Dimensions de l'image affichée: {img_width}x{img_height}")
|
|
|
|
# Ajuster les coordonnées en fonction du zoom et du décalage
|
|
# Note: pour simplifier, nous ignorons le décalage du viewport pour l'instant
|
|
adjusted_rect = QRectF(
|
|
rect.x() / self.zoom_factor / 2, # Division par 2 car la matrice de rendu utilise 2x
|
|
rect.y() / self.zoom_factor / 2,
|
|
rect.width() / self.zoom_factor / 2,
|
|
rect.height() / self.zoom_factor / 2
|
|
)
|
|
|
|
print(f"DEBUG: Coordonnées ajustées: x={adjusted_rect.x()}, y={adjusted_rect.y()}, w={adjusted_rect.width()}, h={adjusted_rect.height()}")
|
|
|
|
selection = {
|
|
"page": page,
|
|
"rect": adjusted_rect,
|
|
"type": "schéma", # Type par défaut
|
|
"context": ""
|
|
}
|
|
|
|
self.selected_regions.append(selection)
|
|
self.config_panel.update_selection_list()
|
|
|
|
# Informer l'utilisateur
|
|
self.status_bar.showMessage(f"Sélection ajoutée à la page {page + 1}")
|
|
|
|
def get_selection_image(self, selection):
|
|
"""
|
|
Extrait l'image de la sélection
|
|
|
|
Args:
|
|
selection (dict): Dictionnaire contenant les informations de la sélection
|
|
|
|
Returns:
|
|
bytes: Données de l'image en bytes, ou None en cas d'erreur
|
|
"""
|
|
if not selection or not self.pdf_document:
|
|
return None
|
|
|
|
try:
|
|
# Récupérer la page
|
|
page_num = selection.get("page", 0)
|
|
page = self.pdf_document[page_num]
|
|
|
|
# Récupérer les coordonnées de la sélection
|
|
rect = selection.get("rect")
|
|
if not rect:
|
|
return None
|
|
|
|
# Convertir les coordonnées en rectangle PyMuPDF
|
|
# Les coordonnées stockées sont déjà ajustées par rapport au zoom (voir méthode add_selection)
|
|
mupdf_rect = fitz.Rect(
|
|
rect.x(),
|
|
rect.y(),
|
|
rect.x() + rect.width(),
|
|
rect.y() + rect.height()
|
|
)
|
|
|
|
# Debug: Afficher les coordonnées pour vérification
|
|
print(f"Coordonnées de la sélection (ajustées): {mupdf_rect}")
|
|
|
|
# Matrice de transformation pour la qualité (sans zoom supplémentaire)
|
|
matrix = fitz.Matrix(2, 2) # Facteur de qualité fixe à 2
|
|
|
|
# Capturer l'image de la zone sélectionnée
|
|
try:
|
|
# Essayer avec la nouvelle API (PyMuPDF récent)
|
|
pix = page.get_pixmap(matrix=matrix, clip=mupdf_rect)
|
|
except AttributeError:
|
|
# Fallback pour les anciennes versions
|
|
pix = page.render_pixmap(matrix=matrix, clip=mupdf_rect)
|
|
|
|
# Debug: Afficher les dimensions du pixmap obtenu
|
|
print(f"Dimensions de l'image capturée: {pix.width}x{pix.height}")
|
|
|
|
# Convertir en format PNG
|
|
img_data = pix.tobytes("png")
|
|
|
|
return img_data
|
|
|
|
except Exception as e:
|
|
print(f"Erreur lors de l'extraction de l'image: {str(e)}")
|
|
return None
|
|
|
|
class PDFLabel(QLabel):
|
|
"""Étiquette personnalisée pour afficher le PDF et gérer les sélections"""
|
|
|
|
selection_made = pyqtSignal(QRectF)
|
|
|
|
def __init__(self, parent):
|
|
super().__init__(parent)
|
|
self.parent = parent
|
|
self.setAlignment(Qt.AlignmentFlag.AlignCenter)
|
|
|
|
# Pour permettre de suivre les événements de souris
|
|
self.setMouseTracking(True)
|
|
|
|
def mousePressEvent(self, event):
|
|
"""Gère l'événement de clic de souris pour démarrer une sélection"""
|
|
if event.button() == Qt.MouseButton.LeftButton:
|
|
self.parent.selection_start = event.position()
|
|
self.parent.selection_active = True
|
|
self.parent.selection_rect = QRectF(event.position(), event.position())
|
|
|
|
def mouseMoveEvent(self, event):
|
|
"""Gère l'événement de déplacement de souris pour mettre à jour la sélection"""
|
|
if self.parent.selection_active:
|
|
# Mise à jour du rectangle de sélection
|
|
self.parent.selection_rect = QRectF(
|
|
self.parent.selection_start,
|
|
event.position()
|
|
).normalized()
|
|
# Demande de mise à jour de l'affichage
|
|
self.update()
|
|
|
|
def mouseReleaseEvent(self, event):
|
|
"""Gère l'événement de relâchement de souris pour finaliser une sélection"""
|
|
if event.button() == Qt.MouseButton.LeftButton and self.parent.selection_active:
|
|
# Finalisation de la sélection
|
|
self.parent.selection_active = False
|
|
|
|
# Vérifier que la sélection a une taille minimale
|
|
if (self.parent.selection_rect.width() > 10 and
|
|
self.parent.selection_rect.height() > 10):
|
|
|
|
# Ajouter cette sélection à la liste des sélections
|
|
self.parent.add_selection(self.parent.selection_rect)
|
|
|
|
# Émettre le signal de sélection
|
|
self.selection_made.emit(self.parent.selection_rect)
|
|
|
|
# Continuer à afficher le rectangle
|
|
self.update()
|
|
|
|
def paintEvent(self, event):
|
|
"""Surcharge pour dessiner la sélection par-dessus le PDF"""
|
|
super().paintEvent(event)
|
|
|
|
# Dessiner le rectangle de sélection si disponible
|
|
if self.parent.selection_rect is not None:
|
|
painter = QPainter(self)
|
|
|
|
# Paramètres pour le rectangle de sélection
|
|
pen = QPen(QColor(255, 0, 0)) # Rouge
|
|
pen.setWidth(2)
|
|
pen.setStyle(Qt.PenStyle.DashLine)
|
|
painter.setPen(pen)
|
|
|
|
# Dessiner avec une légère transparence
|
|
painter.setOpacity(0.7)
|
|
painter.drawRect(self.parent.selection_rect) |