diff --git a/main.py b/main.py index 9d28f31..13c15c8 100644 --- a/main.py +++ b/main.py @@ -1,5 +1,5 @@ import sys -from typing import Optional, Dict, List +from typing import Optional, Dict, List, Tuple from PyQt5.QtWidgets import ( QApplication, QMainWindow, @@ -8,12 +8,14 @@ QAction, QVBoxLayout, QWidget, + QFileDialog, ) from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QMouseEvent, QPaintEvent from PyQt5.QtCore import Qt, QRect, QPoint import fitz # PyMuPDF from PIL import Image import io +import os class DrawableQLabel(QLabel): @@ -22,10 +24,19 @@ def __init__(self) -> None: self.begin: Optional[QPoint] = None self.end: Optional[QPoint] = None self.is_drawing: bool = False - self.current_page_boxes: Dict[ - int, List[QRect] - ] = {} # Dictionary to store boxes for each page + self.current_page_boxes: Dict[int, List[QRect]] = {} self.current_page: int = 0 + self.scale_factor: float = 1.0 + + def set_scale_factor(self, factor: float) -> None: + self.scale_factor = factor + + def get_pdf_coordinates(self, rect: QRect) -> Tuple[float, float, float, float]: + x1 = rect.left() / self.scale_factor + y1 = rect.top() / self.scale_factor + x2 = rect.right() / self.scale_factor + y2 = rect.bottom() / self.scale_factor + return (x1, y1, x2, y2) def set_page(self, page_number: int) -> None: self.current_page = page_number @@ -57,7 +68,6 @@ def mouseReleaseEvent(self, ev: QMouseEvent | None) -> None: if ev.button() == Qt.MouseButton.LeftButton: self.is_drawing = False if self.begin is not None and self.end is not None: - # Add the completed box to the current page's list self.current_page_boxes[self.current_page].append( QRect(self.begin, self.end).normalized() ) @@ -72,12 +82,10 @@ def paintEvent(self, a0: QPaintEvent | None) -> None: painter = QPainter(self) painter.setPen(QPen(Qt.GlobalColor.red, 2, Qt.PenStyle.SolidLine)) - # Draw all saved boxes for the current page if self.current_page in self.current_page_boxes: for box in self.current_page_boxes[self.current_page]: painter.drawRect(box) - # Draw the current box being created if self.begin is not None and self.end is not None and self.is_drawing: painter.drawRect(QRect(self.begin, self.end).normalized()) @@ -100,7 +108,6 @@ def __init__(self, pdf_path: str) -> None: self.setWindowTitle("PDF Viewer") self.setGeometry(100, 100, 800, 600) - # Load PDF self.pdf_path: str = pdf_path self.doc: fitz.Document = fitz.open(pdf_path) self.num_pages: int = len(self.doc) @@ -137,10 +144,20 @@ def __init__(self, pdf_path: str) -> None: self.clear_all_action.triggered.connect(self.clear_all_boxes) self.toolbar.addAction(self.clear_all_action) + # Save Annotations Button + self.save_action: QAction = QAction("Save Annotations", self) + self.save_action.triggered.connect(self.save_annotations) + self.toolbar.addAction(self.save_action) + + # Crop PDF Button + self.crop_action: QAction = QAction("Crop PDF", self) + self.crop_action.triggered.connect(self.crop_pdf_from_annotations) + self.toolbar.addAction(self.crop_action) + # Layout self.main_layout: QVBoxLayout = QVBoxLayout() - # Image Label (now using custom DrawableQLabel) + # Image Label self.image_label: DrawableQLabel = DrawableQLabel() self.image_label.setAlignment(Qt.AlignmentFlag.AlignCenter) self.main_layout.addWidget(self.image_label) @@ -154,24 +171,36 @@ def __init__(self, pdf_path: str) -> None: self.update_page() def update_page(self) -> None: - image: Optional[Image.Image] = self.convert_pdf_page_to_image(self.current_page) - if image: - qimage: QImage = QImage( - image.tobytes(), - image.width, - image.height, - image.width * 3, - QImage.Format.Format_RGB888, - ) - pixmap: QPixmap = QPixmap.fromImage(qimage) - self.image_label.setPixmap(pixmap) - self.image_label.set_page( - self.current_page - ) # Update current page in DrawableQLabel - self.page_label.setText(f"Page {self.current_page + 1} / {self.num_pages}") - self.setWindowTitle( - f"PDF Viewer - Page {self.current_page + 1} / {self.num_pages}" - ) + page = self.doc.load_page(self.current_page) + pix = page.get_pixmap() # type: ignore [attr-defined] + + window_width = self.width() - 50 + scale_factor = window_width / pix.width + + image = Image.open(io.BytesIO(pix.tobytes("png"))) + scaled_size = ( + int(image.width * scale_factor), + int(image.height * scale_factor), + ) + image = image.resize(scaled_size, Image.LANCZOS) # type: ignore [attr-defined] + + qimage = QImage( + image.tobytes(), + image.width, + image.height, + image.width * 3, + QImage.Format.Format_RGB888, + ) + + pixmap = QPixmap.fromImage(qimage) + self.image_label.setPixmap(pixmap) + self.image_label.set_page(self.current_page) + self.image_label.set_scale_factor(scale_factor) + + self.page_label.setText(f"Page {self.current_page + 1} / {self.num_pages}") + self.setWindowTitle( + f"PDF Viewer - Page {self.current_page + 1} / {self.num_pages}" + ) def clear_current_page_boxes(self) -> None: self.image_label.clear_boxes() @@ -179,12 +208,6 @@ def clear_current_page_boxes(self) -> None: def clear_all_boxes(self) -> None: self.image_label.clear_all_boxes() - def convert_pdf_page_to_image(self, page_number: int) -> Optional[Image.Image]: - page: fitz.Page = self.doc.load_page(page_number) # Load specific page - pix: fitz.Pixmap = page.get_pixmap() # type: ignore [attr-defined] - image: Image.Image = Image.open(io.BytesIO(pix.tobytes("png"))) - return image - def show_previous_page(self) -> None: if self.current_page > 0: self.current_page -= 1 @@ -195,10 +218,70 @@ def show_next_page(self) -> None: self.current_page += 1 self.update_page() + def save_annotations(self) -> None: + file_path, _ = QFileDialog.getSaveFileName( + self, "Save Annotations", "", "Markdown Files (*.md)" + ) + if not file_path: + return + + with open(file_path, "w", encoding="utf-8") as f: + f.write("# PDF Annotations\n\n") + boxes = self.image_label.get_boxes() + for page_num in sorted(boxes.keys()): + if boxes[page_num]: # 박스가 있는 페이지만 저장 + f.write(f"## Page {page_num + 1}\n\n") + for i, rect in enumerate(boxes[page_num], 1): + x1, y1, x2, y2 = self.image_label.get_pdf_coordinates(rect) + f.write(f"### Box {i}\n") + f.write( + f"- Coordinates: [{x1:.2f}, {y1:.2f}, {x2:.2f}, {y2:.2f}]\n\n" + ) + + def crop_pdf_from_annotations(self) -> None: + md_path, _ = QFileDialog.getOpenFileName( + self, "Open Annotations", "", "Markdown Files (*.md)" + ) + if not md_path: + return + + output_dir, _ = QFileDialog.getSaveFileName( + self, "Save Cropped Images", "", "Select Directory" + ) + if not output_dir: + return + + output_dir = os.path.dirname(output_dir) + os.makedirs(output_dir, exist_ok=True) + + current_page = None + with open(md_path, "r", encoding="utf-8") as f: + for line in f: + if line.startswith("## Page"): + current_page = int(line.split()[2]) - 1 + elif line.startswith("- Coordinates:"): + coords = eval(line.split(": ")[1].strip()) + if current_page is not None: + self.crop_and_save_region(current_page, coords, output_dir) + + def crop_and_save_region( + self, page_num: int, coords: List[float], output_dir: str + ) -> None: + page = self.doc.load_page(page_num) + x1, y1, x2, y2 = coords + + rect = fitz.Rect(x1, y1, x2, y2) + pix = page.get_pixmap(clip=rect) # type: ignore [attr-defined] + + image_path = os.path.join( + output_dir, f"page_{page_num + 1}_crop_{x1:.0f}_{y1:.0f}.png" + ) + pix.save(image_path) + if __name__ == "__main__": app: QApplication = QApplication(sys.argv) - pdf_path: str = "crawler.pdf" # Replace with your PDF file path + pdf_path: str = "crawler.pdf" viewer: PDFViewer = PDFViewer(pdf_path) viewer.show() sys.exit(app.exec_())