Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. Weโ€™ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

โœจ feat(main) : add export by md #7

Merged
merged 2 commits into from
Dec 29, 2024
Merged

โœจ feat(main) : add export by md #7

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
8fb04ae
โœจ feat(main) : add export by md
Tagite Dec 29, 2024
89aa052
๐ŸŽจ style(main) : change code format by ruff
Tagite Dec 29, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
151 changes: 117 additions & 34 deletions main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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):
Expand All @@ -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
Expand Down Expand Up @@ -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()
)
Expand All @@ -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())

Expand All @@ -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)
Expand Down Expand Up @@ -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)
Expand All @@ -154,37 +171,43 @@ 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()

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
Expand All @@ -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_())
Loading