296 lines
11 KiB
Python
296 lines
11 KiB
Python
import os
|
|
import fitz # PyMuPDF
|
|
import torch
|
|
from PIL import Image
|
|
import numpy as np
|
|
import folder_paths
|
|
from .utils import strip_path
|
|
from docx import Document
|
|
|
|
class DocumentLoaderNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith((".pdf", ".txt", ".doc", ".docx"))]
|
|
return {"required": {"file_path": (sorted(files),)}}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("parsed_text",)
|
|
FUNCTION = "load_document"
|
|
CATEGORY = "document_processing"
|
|
|
|
def load_document(self, file_path):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
_, ext = os.path.splitext(full_path)
|
|
if ext.lower() == '.pdf':
|
|
return (self.parse_pdf(full_path),)
|
|
elif ext.lower() in ['.txt', '.doc', '.docx']:
|
|
return (self.parse_text(full_path),)
|
|
else:
|
|
raise ValueError(f"Unsupported file type: {ext}")
|
|
|
|
def parse_pdf(self, file_path):
|
|
text = ""
|
|
with fitz.open(file_path) as doc:
|
|
for page in doc:
|
|
text += page.get_text()
|
|
return text
|
|
|
|
def parse_text(self, file_path):
|
|
_, ext = os.path.splitext(file_path)
|
|
if ext.lower() == '.txt':
|
|
with open(file_path, 'r', encoding='utf-8') as file:
|
|
return file.read()
|
|
elif ext.lower() in ['.doc', '.docx']:
|
|
doc = Document(file_path)
|
|
return "\n".join([para.text for para in doc.paragraphs])
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, file_path):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
return full_path
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, file_path):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
return os.path.isfile(full_path)
|
|
|
|
class PDFToImageNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith(".pdf")]
|
|
return {
|
|
"required": {
|
|
"file_path": (sorted(files),),
|
|
"start_page": ("INT", {"default": 1, "min": 1, "max": 10000}),
|
|
"end_page": ("INT", {"default": 1, "min": 1, "max": 10000}),
|
|
"dpi": ("INT", {"default": 300, "min": 72, "max": 600}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "pdf_to_images"
|
|
CATEGORY = "document_processing"
|
|
OUTPUT_IS_LIST = (True,)
|
|
|
|
def pdf_to_images(self, file_path, start_page, end_page, dpi):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
doc = fitz.open(full_path)
|
|
num_pages = len(doc)
|
|
|
|
start_page = max(1, min(start_page, num_pages))
|
|
end_page = max(start_page, min(end_page, num_pages))
|
|
|
|
images = []
|
|
for page_num in range(start_page - 1, end_page):
|
|
page = doc.load_page(page_num)
|
|
pix = page.get_pixmap(matrix=fitz.Matrix(dpi/72, dpi/72))
|
|
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
|
img_np = np.array(img).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_np)[None,]
|
|
images.append(img_tensor)
|
|
|
|
doc.close()
|
|
return (images,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, file_path, start_page, end_page, dpi):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
return full_path
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, file_path, start_page, end_page, dpi):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
return os.path.isfile(full_path)
|
|
|
|
class PDFPageSplitterNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
input_dir = folder_paths.get_input_directory()
|
|
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f)) and f.lower().endswith(".pdf")]
|
|
return {
|
|
"required": {
|
|
"file_path": (sorted(files),),
|
|
"page_numbers": ("STRING", {"default": "1,2,3"}),
|
|
"dpi": ("INT", {"default": 300, "min": 72, "max": 600}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
RETURN_NAMES = ("selected_pages",)
|
|
FUNCTION = "pdf_to_selected_images"
|
|
CATEGORY = "document_processing"
|
|
OUTPUT_NODE = True
|
|
OUTPUT_IS_LIST = (True,)
|
|
|
|
def pdf_to_selected_images(self, file_path, page_numbers, dpi):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
doc = fitz.open(full_path)
|
|
num_pages = len(doc)
|
|
|
|
# Parse page numbers from input string
|
|
try:
|
|
selected_pages = [int(p.strip()) for p in page_numbers.split(',') if p.strip()]
|
|
selected_pages = [p for p in selected_pages if 1 <= p <= num_pages]
|
|
except ValueError:
|
|
raise ValueError("Invalid page numbers. Please provide comma-separated integers.")
|
|
|
|
images = []
|
|
for page_num in selected_pages:
|
|
page = doc.load_page(page_num - 1) # PDF pages are 0-indexed
|
|
pix = page.get_pixmap(matrix=fitz.Matrix(dpi/72, dpi/72))
|
|
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
|
|
img_np = np.array(img).astype(np.float32) / 255.0
|
|
img_tensor = torch.from_numpy(img_np)[None,]
|
|
images.append(img_tensor)
|
|
|
|
doc.close()
|
|
return (images,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, file_path, page_numbers, dpi):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
return full_path
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(cls, file_path, page_numbers, dpi):
|
|
full_path = folder_paths.get_annotated_filepath(strip_path(file_path))
|
|
if not os.path.isfile(full_path):
|
|
return False
|
|
try:
|
|
[int(p.strip()) for p in page_numbers.split(',') if p.strip()]
|
|
except ValueError:
|
|
return False
|
|
return True
|
|
|
|
class ImageSelectorNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"indexes": ("STRING", {"default": "0"}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("IMAGE",)
|
|
FUNCTION = "select_images"
|
|
CATEGORY = "image_processing"
|
|
OUTPUT_NODE = True
|
|
OUTPUT_IS_LIST = (True,)
|
|
|
|
def select_images(self, images, indexes):
|
|
try:
|
|
index_list = [int(idx.strip()) for idx in indexes.split(',')]
|
|
except ValueError:
|
|
raise ValueError("Invalid indexes. Please provide comma-separated integers.")
|
|
|
|
selected_images = []
|
|
for idx in index_list:
|
|
if idx < 0 or idx >= len(images):
|
|
raise ValueError(f"Index {idx} is out of range. Available images: {len(images)}")
|
|
selected_images.append(images[idx])
|
|
|
|
return (selected_images,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, images, indexes):
|
|
return indexes
|
|
|
|
class TextChunkerNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"text": ("STRING", {"multiline": True}),
|
|
"chunk_size": ("INT", {"default": 1000, "min": 1, "max": 10000}),
|
|
"chunk_method": (["words", "characters"],),
|
|
"respect_word_boundaries": ("BOOLEAN", {"default": True}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("text_chunks",)
|
|
FUNCTION = "chunk_text"
|
|
CATEGORY = "document_processing"
|
|
OUTPUT_NODE = True
|
|
OUTPUT_IS_LIST = (True,)
|
|
|
|
def chunk_text(self, text, chunk_size, chunk_method, respect_word_boundaries):
|
|
if not text:
|
|
return (["No text provided"],)
|
|
|
|
chunks = []
|
|
if chunk_method == "words":
|
|
words = text.split()
|
|
current_chunk = []
|
|
word_count = 0
|
|
for word in words:
|
|
current_chunk.append(word)
|
|
word_count += 1
|
|
if word_count >= chunk_size:
|
|
chunks.append(" ".join(current_chunk))
|
|
current_chunk = []
|
|
word_count = 0
|
|
if current_chunk:
|
|
chunks.append(" ".join(current_chunk))
|
|
else: # characters
|
|
if respect_word_boundaries:
|
|
words = text.split()
|
|
current_chunk = []
|
|
char_count = 0
|
|
for word in words:
|
|
if char_count + len(word) > chunk_size and current_chunk:
|
|
chunks.append(" ".join(current_chunk))
|
|
current_chunk = []
|
|
char_count = 0
|
|
current_chunk.append(word)
|
|
char_count += len(word) + 1 # +1 for space
|
|
if current_chunk:
|
|
chunks.append(" ".join(current_chunk))
|
|
else:
|
|
chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]
|
|
|
|
return (chunks,)
|
|
|
|
class ChunkRouterNode:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"chunks": ("STRING", {"forceInput": True}),
|
|
"indices": ("INT", {"forceInput": True}),
|
|
"selected_index": ("INT", {"default": 0, "min": 0, "max": 1000}),
|
|
}
|
|
}
|
|
|
|
RETURN_TYPES = ("STRING",)
|
|
RETURN_NAMES = ("selected_chunk",)
|
|
FUNCTION = "route_chunk"
|
|
CATEGORY = "document_processing"
|
|
|
|
def route_chunk(self, chunks, indices, selected_index):
|
|
if selected_index < 0 or selected_index >= len(chunks):
|
|
raise ValueError(f"Selected index {selected_index} is out of range. Available chunks: {len(chunks)}")
|
|
return (chunks[selected_index],)
|
|
|
|
# Update NODE_CLASS_MAPPINGS and NODE_DISPLAY_NAME_MAPPINGS
|
|
NODE_CLASS_MAPPINGS = {
|
|
"DocumentLoader": DocumentLoaderNode,
|
|
"PDFToImage": PDFToImageNode,
|
|
"PDFPageSplitter": PDFPageSplitterNode,
|
|
"ImageSelector": ImageSelectorNode,
|
|
"TextChunker": TextChunkerNode,
|
|
"ChunkRouter": ChunkRouterNode
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"DocumentLoader": "Document Loader",
|
|
"PDFToImage": "PDF to Image (Multi-Page)",
|
|
"PDFPageSplitter": "PDF Page Splitter",
|
|
"ImageSelector": "Image Selector",
|
|
"TextChunker": "Text Chunker",
|
|
"ChunkRouter": "Chunk Router"
|
|
}
|