import os import numpy as np import cv2 from PIL import Image, ImageDraw, ImageFont import torch import folder_paths from typing import List, Dict, Any, Tuple from comfy.utils import common_upscale try: from comfy.utils import PIL_to_tensor, tensor_to_PIL except ImportError: # Fallback to manual implementation import torchvision.transforms.functional as F def PIL_to_tensor(pil_image): return F.to_tensor(pil_image).permute(1, 2, 0) def tensor_to_PIL(tensor): if tensor.dim() == 4: tensor = tensor[0] return F.to_pil_image(tensor.permute(2, 0, 1)) # LayoutLMv3 and transformers imports try: from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification from transformers import LayoutLMv3ForSequenceClassification, LayoutLMv3ForQuestionAnswering LAYOUTLMV3_AVAILABLE = True except ImportError as e: LAYOUTLMV3_AVAILABLE = False print("LayoutLMv3 not available:", e) print("Install with: pip install transformers torch torchvision") # OCR imports - using EasyOCR as primary choice for simplicity try: import easyocr EASYOCR_AVAILABLE = True except ImportError: EASYOCR_AVAILABLE = False print("EasyOCR not available. Install with: pip install easyocr") try: from paddleocr import PaddleOCR PADDLEOCR_AVAILABLE = True except ImportError: PADDLEOCR_AVAILABLE = False print("PaddleOCR not available. Install with: pip install paddleocr") class LayoutLMv3DocumentAnalysisNode: """Advanced document analysis using Microsoft's LayoutLMv3""" @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "model_size": (["base", "large"], {"default": "base"}), "task_type": (["token_classification", "sequence_classification"], {"default": "token_classification"}), "ocr_engine": (["easyocr", "paddleocr"], {"default": "easyocr"}), "confidence_threshold": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0}), "extract_crops": ("BOOLEAN", {"default": True}), "language": (["en", "ch", "fr", "de", "ja", "ko"], {"default": "en"}), }, "optional": { "custom_labels": ("STRING", {"default": "", "multiline": True}), } } RETURN_TYPES = ("IMAGE", "LIST", "LIST", "LIST", "LIST", "DICT") RETURN_NAMES = ("overlay_image", "image_boxes", "text_boxes", "image_crops", "text_tokens", "analysis_results") FUNCTION = "run" CATEGORY = "layout/advanced" def __init__(self): if not LAYOUTLMV3_AVAILABLE: raise ImportError("LayoutLMv3 requires transformers. Install with: pip install transformers") self.processor = None self.model = None self.ocr_reader = None self.current_model = None self.current_task = None # Default label mappings for document layout analysis self.layout_labels = { 0: "O", # Outside/Other 1: "B-TITLE", # Beginning of title 2: "I-TITLE", # Inside title 3: "B-HEADER", # Beginning of header 4: "I-HEADER", # Inside header 5: "B-TEXT", # Beginning of text block 6: "I-TEXT", # Inside text block 7: "B-LIST", # Beginning of list 8: "I-LIST", # Inside list 9: "B-TABLE", # Beginning of table 10: "I-TABLE", # Inside table 11: "B-FIGURE", # Beginning of figure 12: "I-FIGURE", # Inside figure } def _init_models(self, model_size="base", task_type="token_classification", ocr_engine="easyocr", language="en"): """Initialize LayoutLMv3 models and OCR engine""" # Initialize OCR engine if ocr_engine == "easyocr" and EASYOCR_AVAILABLE: if self.ocr_reader is None: self.ocr_reader = easyocr.Reader([language], gpu=torch.cuda.is_available()) elif ocr_engine == "paddleocr" and PADDLEOCR_AVAILABLE: if self.ocr_reader is None: self.ocr_reader = PaddleOCR( use_angle_cls=True, lang=language, use_gpu=torch.cuda.is_available(), show_log=False ) else: raise ImportError(f"OCR engine {ocr_engine} not available") # Initialize LayoutLMv3 model if needed model_name = f"microsoft/layoutlmv3-{model_size}" if self.current_model != model_name or self.current_task != task_type: print(f"Loading LayoutLMv3 model: {model_name} for {task_type}") # Initialize processor with apply_ocr=False since we're doing external OCR self.processor = LayoutLMv3Processor.from_pretrained(model_name, apply_ocr=False) if task_type == "token_classification": self.model = LayoutLMv3ForTokenClassification.from_pretrained(model_name) elif task_type == "sequence_classification": self.model = LayoutLMv3ForSequenceClassification.from_pretrained(model_name) # Move to GPU if available if torch.cuda.is_available(): self.model = self.model.cuda() self.model.eval() self.current_model = model_name self.current_task = task_type def _run_ocr(self, image_array: np.ndarray, ocr_engine: str) -> List[Dict]: """Run OCR to extract text and bounding boxes""" raw_results = [] if ocr_engine == "easyocr": ocr_results = self.ocr_reader.readtext(image_array) for bbox, text, confidence in ocr_results: # Convert bbox format x_coords = [p[0] for p in bbox] y_coords = [p[1] for p in bbox] x1, y1, x2, y2 = min(x_coords), min(y_coords), max(x_coords), max(y_coords) raw_results.append({ 'text': text, 'bbox': [int(x1), int(y1), int(x2), int(y2)], 'confidence': confidence }) elif ocr_engine == "paddleocr": ocr_results = self.ocr_reader.ocr(image_array, cls=True) if ocr_results and ocr_results[0]: for line in ocr_results[0]: bbox = line[0] text = line[1][0] confidence = line[1][1] # Convert bbox format x_coords = [p[0] for p in bbox] y_coords = [p[1] for p in bbox] x1, y1, x2, y2 = min(x_coords), min(y_coords), max(x_coords), max(y_coords) raw_results.append({ 'text': text, 'bbox': [int(x1), int(y1), int(x2), int(y2)], 'confidence': confidence }) # Group nearby text detections into text blocks grouped_results = self._group_text_into_blocks(raw_results) print(f"Grouped {len(raw_results)} text detections into {len(grouped_results)} text blocks") return grouped_results def _group_text_into_blocks(self, text_detections: List[Dict]) -> List[Dict]: """Group nearby text detections into logical text blocks""" if not text_detections: return [] # Sort by vertical position (top to bottom) text_detections.sort(key=lambda x: x['bbox'][1]) blocks = [] current_block = None for detection in text_detections: x1, y1, x2, y2 = detection['bbox'] if current_block is None: # Start first block current_block = { 'text': detection['text'], 'bbox': [x1, y1, x2, y2], 'confidence': detection['confidence'], 'detections': [detection] } else: # Check if this detection should be merged with current block curr_x1, curr_y1, curr_x2, curr_y2 = current_block['bbox'] # Calculate vertical and horizontal proximity vertical_gap = y1 - curr_y2 horizontal_overlap = min(x2, curr_x2) - max(x1, curr_x1) # Merge conditions: # 1. Small vertical gap (same paragraph) # 2. Some horizontal overlap or close proximity should_merge = ( vertical_gap < 30 and # Lines close vertically (horizontal_overlap > -50 or # Some horizontal overlap abs(x1 - curr_x1) < 100) # Or similar left alignment ) if should_merge: # Merge with current block current_block['text'] += ' ' + detection['text'] current_block['bbox'] = [ min(curr_x1, x1), # leftmost min(curr_y1, y1), # topmost max(curr_x2, x2), # rightmost max(curr_y2, y2) # bottommost ] current_block['confidence'] = min(current_block['confidence'], detection['confidence']) current_block['detections'].append(detection) else: # Start new block blocks.append(current_block) current_block = { 'text': detection['text'], 'bbox': [x1, y1, x2, y2], 'confidence': detection['confidence'], 'detections': [detection] } # Add the last block if current_block: blocks.append(current_block) return blocks def _prepare_layoutlmv3_inputs(self, image: Image.Image, ocr_results: List[Dict]) -> Dict: """Prepare inputs for LayoutLMv3 with aggressive validation""" # Extract text and bounding boxes words = [] boxes = [] # Limit OCR results to prevent memory issues max_regions = 50 # Limit to 50 regions to prevent crashes ocr_results = ocr_results[:max_regions] img_width, img_height = image.size print(f"🔍 Image dimensions: {img_width}x{img_height}") for result in ocr_results: # Split text into words and assign same bbox to each word text_words = result['text'].split() for word in text_words: # Skip empty words if not word.strip(): continue words.append(word) # Aggressive bounding box validation and fixing x1, y1, x2, y2 = result['bbox'] # Convert to integers and validate x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2) # Ensure positive coordinates x1 = max(0, x1) y1 = max(0, y1) x2 = max(0, x2) y2 = max(0, y2) # Fix invalid bounding boxes if x1 >= x2: x2 = x1 + 10 # Minimum width if y1 >= y2: y2 = y1 + 10 # Minimum height # Clamp to image bounds with safety margin x1 = min(x1, img_width - 2) y1 = min(y1, img_height - 2) x2 = min(x2, img_width - 1) y2 = min(y2, img_height - 1) # Final validation if x1 >= x2: x2 = x1 + 1 if y1 >= y2: y2 = y1 + 1 # Ensure reasonable size (not too small) if (x2 - x1) < 1: x2 = x1 + 1 if (y2 - y1) < 1: y2 = y1 + 1 boxes.append([x1, y1, x2, y2]) # Limit total words to LayoutLMv3's max sequence length max_words = 300 # More conservative limit if len(words) > max_words: print(f"Warning: Truncating from {len(words)} to {max_words} words") words = words[:max_words] boxes = boxes[:max_words] print(f"Processing {len(words)} words from {len(ocr_results)} regions") # Validate all boxes one more time valid_words = [] valid_boxes = [] for word, box in zip(words, boxes): x1, y1, x2, y2 = box # Final sanity checks if (x1 < x2 and y1 < y2 and x1 >= 0 and y1 >= 0 and x2 <= img_width and y2 <= img_height and len(word.strip()) > 0): valid_words.append(word) valid_boxes.append(box) else: print(f"⚠️ Skipping invalid word/box: '{word}' {box}") print(f"After validation: {len(valid_words)} valid words") if not valid_words: raise ValueError("No valid words/boxes after validation") try: # Prepare inputs using LayoutLMv3 processor encoding = self.processor( image, valid_words, boxes=valid_boxes, return_tensors="pt", truncation=True, padding=True, max_length=256 # Even more conservative ) # Move to GPU if available if torch.cuda.is_available(): encoding = {k: v.cuda() for k, v in encoding.items()} return encoding, valid_words, valid_boxes except Exception as e: print(f"Error in processor: {e}") raise e def _process_token_classification_results(self, predictions, words, boxes, confidence_threshold): """Process token classification results into structured layout information""" text_boxes = [] image_boxes = [] processed_tokens = [] # Get predictions predicted_labels = torch.argmax(predictions.logits, dim=-1) confidences = torch.softmax(predictions.logits, dim=-1) # Process each token for i, (word, box, label_id, token_confidences) in enumerate(zip(words, boxes, predicted_labels[0], confidences[0])): label_id = label_id.item() confidence = token_confidences[label_id].item() if confidence < confidence_threshold: continue # Map label ID to label name label_name = self.layout_labels.get(label_id, f"LABEL_{label_id}") token_info = { "word": word, "label": label_name, "confidence": confidence, "bbox": box } processed_tokens.append(token_info) # Categorize into text or image elements if any(keyword in label_name.lower() for keyword in ['title', 'header', 'text', 'list']): text_boxes.append({ "type": label_name.replace('B-', '').replace('I-', '').lower(), "score": confidence, "x1": box[0], "y1": box[1], "x2": box[2], "y2": box[3], "text": word }) elif any(keyword in label_name.lower() for keyword in ['table', 'figure']): image_boxes.append({ "type": label_name.replace('B-', '').replace('I-', '').lower(), "score": confidence, "x1": box[0], "y1": box[1], "x2": box[2], "y2": box[3], "text": word }) return text_boxes, image_boxes, processed_tokens def _create_overlay_visualization(self, image: Image.Image, text_boxes: List, image_boxes: List): """Create visualization overlay""" overlay = image.copy() draw = ImageDraw.Draw(overlay) # Color mapping colors = { "title": (255, 0, 0), # Red "header": (255, 165, 0), # Orange "text": (0, 255, 0), # Green "list": (0, 255, 255), # Cyan "table": (0, 0, 255), # Blue "figure": (255, 255, 0), # Yellow } # Draw text boxes for box in text_boxes: color = colors.get(box["type"], (128, 128, 128)) draw.rectangle([box["x1"], box["y1"], box["x2"], box["y2"]], outline=color, width=2) # Add label label = f'{box["type"]}: {box["score"]:.2f}' draw.text((box["x1"], max(0, box["y1"]-15)), label, fill=color) # Draw image boxes for box in image_boxes: color = colors.get(box["type"], (128, 128, 128)) draw.rectangle([box["x1"], box["y1"], box["x2"], box["y2"]], outline=color, width=3) # Add label label = f'{box["type"]}: {box["score"]:.2f}' draw.text((box["x1"], max(0, box["y1"]-15)), label, fill=color) return overlay def _debug_bounding_boxes(self, words, boxes, image_size): """Debug function to investigate invalid bounding boxes""" print("🔍 Debugging bounding boxes...") img_width, img_height = image_size invalid_count = 0 for i, (word, box) in enumerate(zip(words, boxes)): x1, y1, x2, y2 = box issues = [] if x1 >= x2: issues.append(f"x1({x1}) >= x2({x2})") if y1 >= y2: issues.append(f"y1({y1}) >= y2({y2})") if x1 < 0 or y1 < 0: issues.append(f"negative coords: ({x1}, {y1})") if x2 > img_width or y2 > img_height: issues.append(f"exceeds image bounds: ({x2}, {y2}) > ({img_width}, {img_height})") if issues: invalid_count += 1 print(f" Box {i} for word '{word}': {box} - Issues: {', '.join(issues)}") if invalid_count >= 10: # Limit output print(" ... (showing first 10 invalid boxes)") break print(f"🔍 Found {invalid_count} invalid boxes out of {len(boxes)} total") return invalid_count def run(self, image, model_size="base", task_type="token_classification", ocr_engine="easyocr", confidence_threshold=0.5, extract_crops=True, language="en", custom_labels=""): # Initialize models self._init_models(model_size, task_type, ocr_engine, language) # Memory warning for large model if model_size == "large": print("⚠️ Using large model - this requires significant GPU memory!") print("💡 If you get crashes, try using 'base' model instead") # Convert tensor to PIL and RESIZE to prevent coordinate issues pil_image = tensor_to_PIL(image) # RESIZE IMAGE to prevent EasyOCR coordinate issues original_size = pil_image.size if original_size[0] > 1000 or original_size[1] > 1000: # Resize large images to prevent coordinate issues aspect_ratio = original_size[0] / original_size[1] if aspect_ratio > 1: new_size = (1000, int(1000 / aspect_ratio)) else: new_size = (int(1000 * aspect_ratio), 1000) pil_image = pil_image.resize(new_size, Image.Resampling.LANCZOS) print(f"🔧 Resized image from {original_size} to {pil_image.size} to prevent OCR coordinate issues") image_array = np.array(pil_image) # Run OCR to get text and bounding boxes print("Running OCR...") ocr_results = self._run_ocr(image_array, ocr_engine) if not ocr_results: print("No text detected by OCR") # Return empty results overlay_tensor = PIL_to_tensor(pil_image) return (overlay_tensor, [], [], [], [], {"message": "No text detected"}) print(f"OCR detected {len(ocr_results)} text regions") # Prepare LayoutLMv3 inputs print("Preparing LayoutLMv3 inputs...") try: encoding, words, boxes = self._prepare_layoutlmv3_inputs(pil_image, ocr_results) except Exception as e: print(f"❌ Error preparing inputs: {e}") return (PIL_to_tensor(pil_image), [], [], [], [], {"error": f"Input preparation failed: {e}"}) # Run LayoutLMv3 inference print("Running LayoutLMv3 inference...") # Check GPU memory before inference if torch.cuda.is_available(): gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9 memory_used = torch.cuda.memory_allocated() / 1e9 memory_free = gpu_memory - memory_used print(f"GPU Memory: {memory_free:.1f}GB free of {gpu_memory:.1f}GB total") # Initialize outputs variable in the correct scope outputs = None try: # Clear GPU cache before inference if torch.cuda.is_available(): torch.cuda.empty_cache() print("🔍 Debug: About to call model inference...") print(f"🔍 Input shapes:") for key, value in encoding.items(): if hasattr(value, 'shape'): print(f" {key}: {value.shape}") # Check for invalid bbox values and fix them properly if 'bbox' in encoding: bbox_tensor = encoding['bbox'] print(f"🔍 Bbox stats:") print(f" Min: {bbox_tensor.min().item()}") print(f" Max: {bbox_tensor.max().item()}") print(f" Contains NaN: {torch.isnan(bbox_tensor).any().item()}") print(f" Contains Inf: {torch.isinf(bbox_tensor).any().item()}") # Check for invalid bounding boxes in the processed tensor invalid_boxes = torch.logical_or( bbox_tensor[:, :, 0] >= bbox_tensor[:, :, 2], # x1 >= x2 bbox_tensor[:, :, 1] >= bbox_tensor[:, :, 3] # y1 >= y2 ) if invalid_boxes.any(): invalid_count = invalid_boxes.sum().item() print(f"⚠️ Found {invalid_count} invalid bounding boxes in processed tensor!") # Debug the source of invalid boxes print("🔍 Investigating source of invalid boxes...") self._debug_bounding_boxes(words, boxes, pil_image.size) # Fix them in the tensor directly print("🔧 Fixing invalid boxes in processed tensor...") bbox_fixed = bbox_tensor.clone() # Fix x coordinates - ensure x2 > x1 bad_x = bbox_fixed[:, :, 0] >= bbox_fixed[:, :, 2] bbox_fixed[:, :, 2][bad_x] = bbox_fixed[:, :, 0][bad_x] + 1 # Fix y coordinates - ensure y2 > y1 bad_y = bbox_fixed[:, :, 1] >= bbox_fixed[:, :, 3] bbox_fixed[:, :, 3][bad_y] = bbox_fixed[:, :, 1][bad_y] + 1 # Update the encoding encoding['bbox'] = bbox_fixed print("✅ Fixed invalid bounding boxes in tensor!") else: print("✅ All bounding boxes valid in processed tensor!") # Direct inference - we know it works up to 500 tokens print("🔍 Running direct inference...") with torch.no_grad(): outputs = self.model(**encoding) print("✅ Inference successful!") except RuntimeError as e: error_msg = str(e).lower() print(f"❌ Runtime error during inference: {e}") # Check if it's a memory error if "out of memory" in error_msg or "cuda" in error_msg: print("🔧 Trying memory-saving strategies...") # Strategy 1: Remove pixel_values to save memory try: print("🔍 Trying without pixel_values...") fallback_encoding = {k: v for k, v in encoding.items() if k != 'pixel_values'} with torch.no_grad(): outputs = self.model(**fallback_encoding) print("✅ Works without pixel_values!") except Exception as e2: print(f"❌ Still fails without pixel_values: {e2}") # Strategy 2: Try on CPU print("🔍 Trying on CPU...") try: cpu_encoding = {k: v.cpu() if hasattr(v, 'cpu') else v for k, v in encoding.items()} cpu_model = self.model.cpu() with torch.no_grad(): outputs = cpu_model(**cpu_encoding) print("✅ Works on CPU!") # Move model back to GPU for next time if torch.cuda.is_available(): self.model = self.model.cuda() except Exception as e3: print(f"❌ Even CPU fails: {e3}") if torch.cuda.is_available(): torch.cuda.empty_cache() return (PIL_to_tensor(pil_image), [], [], [], [], {"error": f"All inference strategies failed: {str(e)}"}) else: # Non-memory error, return immediately print(f"❌ Non-memory error: {e}") return (PIL_to_tensor(pil_image), [], [], [], [], {"error": f"Inference error: {str(e)}"}) # Clear GPU cache after fallback attempts if torch.cuda.is_available(): torch.cuda.empty_cache() except Exception as e: print(f"❌ Unexpected error during inference: {e}") import traceback traceback.print_exc() return (PIL_to_tensor(pil_image), [], [], [], [], {"error": f"Unexpected error: {str(e)}"}) # Ensure we have outputs before proceeding if outputs is None: return (PIL_to_tensor(pil_image), [], [], [], [], {"error": "No outputs generated from model"}) # Process results based on task type try: if task_type == "token_classification": text_boxes, image_boxes, processed_tokens = self._process_token_classification_results( outputs, words, boxes, confidence_threshold ) # Create analysis results for token classification analysis_results = { "task_type": task_type, "model_used": f"microsoft/layoutlmv3-{model_size}", "total_text_regions": len(text_boxes), "total_image_regions": len(image_boxes), "total_tokens": len(processed_tokens), "ocr_engine": ocr_engine, "language": language } else: # For sequence classification, create simple results predicted_class = torch.argmax(outputs.logits, dim=-1).item() confidence = torch.softmax(outputs.logits, dim=-1).max().item() text_boxes = [] image_boxes = [] processed_tokens = [] # Create analysis results for sequence classification analysis_results = { "task_type": task_type, "model_used": f"microsoft/layoutlmv3-{model_size}", "predicted_class": predicted_class, "confidence": confidence, "total_tokens": len(words), "ocr_engine": ocr_engine, "language": language } except Exception as e: print(f"❌ Error processing results: {e}") return (PIL_to_tensor(pil_image), [], [], [], [], {"error": f"Result processing failed: {str(e)}"}) # Extract crops if requested image_crops = [] if extract_crops: for box_list in [text_boxes, image_boxes]: for box in box_list: try: # Ensure coordinates are within image bounds x1 = max(0, min(box["x1"], pil_image.width - 1)) y1 = max(0, min(box["y1"], pil_image.height - 1)) x2 = max(x1 + 1, min(box["x2"], pil_image.width)) y2 = max(y1 + 1, min(box["y2"], pil_image.height)) crop = pil_image.crop((x1, y1, x2, y2)) image_crops.append(PIL_to_tensor(crop)) except Exception as e: print(f"Error extracting crop: {e}") # Create visualization overlay try: overlay = self._create_overlay_visualization(pil_image, text_boxes, image_boxes) overlay_tensor = PIL_to_tensor(overlay) except Exception as e: print(f"Error creating overlay: {e}") overlay_tensor = PIL_to_tensor(pil_image) return (overlay_tensor, image_boxes, text_boxes, image_crops, processed_tokens, analysis_results) class LayoutLMv3QuestionAnsweringNode: """Document Question Answering using LayoutLMv3""" @classmethod def INPUT_TYPES(cls): return { "required": { "image": ("IMAGE",), "question": ("STRING", {"default": "What is the total amount?"}), "model_size": (["base", "large"], {"default": "base"}), "ocr_engine": (["easyocr", "paddleocr"], {"default": "easyocr"}), "language": (["en", "ch", "fr", "de", "ja", "ko"], {"default": "en"}), } } RETURN_TYPES = ("STRING", "FLOAT", "DICT") RETURN_NAMES = ("answer", "confidence", "details") FUNCTION = "run" CATEGORY = "layout/advanced" def __init__(self): if not LAYOUTLMV3_AVAILABLE: raise ImportError("LayoutLMv3 requires transformers. Install with: pip install transformers") self.processor = None self.model = None self.ocr_reader = None self.current_model = None def _init_models(self, model_size="base", ocr_engine="easyocr", language="en"): """Initialize models""" # Initialize OCR if ocr_engine == "easyocr" and EASYOCR_AVAILABLE: if self.ocr_reader is None: self.ocr_reader = easyocr.Reader([language], gpu=torch.cuda.is_available()) elif ocr_engine == "paddleocr" and PADDLEOCR_AVAILABLE: if self.ocr_reader is None: self.ocr_reader = PaddleOCR(use_angle_cls=True, lang=language, use_gpu=torch.cuda.is_available(), show_log=False) # Initialize LayoutLMv3 for QA model_name = f"microsoft/layoutlmv3-{model_size}" if self.current_model != model_name: print(f"Loading LayoutLMv3 QA model: {model_name}") self.processor = LayoutLMv3Processor.from_pretrained(model_name, apply_ocr=False) self.model = LayoutLMv3ForQuestionAnswering.from_pretrained(model_name) if torch.cuda.is_available(): self.model = self.model.cuda() self.model.eval() self.current_model = model_name def run(self, image, question, model_size="base", ocr_engine="easyocr", language="en"): # Initialize models self._init_models(model_size, ocr_engine, language) # Convert and run OCR (reuse OCR logic from main node) pil_image = tensor_to_PIL(image) image_array = np.array(pil_image) # Simple OCR extraction if ocr_engine == "easyocr": ocr_results = self.ocr_reader.readtext(image_array) words = [] boxes = [] for bbox, text, confidence in ocr_results: text_words = text.split() for word in text_words: words.append(word) # Convert bbox x_coords = [p[0] for p in bbox] y_coords = [p[1] for p in bbox] boxes.append([min(x_coords), min(y_coords), max(x_coords), max(y_coords)]) if not words: return ("No text found", 0.0, {"error": "No text detected by OCR"}) # Prepare inputs for QA encoding = self.processor( pil_image, question, words, boxes=boxes, return_tensors="pt", truncation=True, padding=True ) if torch.cuda.is_available(): encoding = {k: v.cuda() for k, v in encoding.items()} # Run inference with torch.no_grad(): outputs = self.model(**encoding) # Extract answer start_scores = outputs.start_logits end_scores = outputs.end_logits start_idx = torch.argmax(start_scores) end_idx = torch.argmax(end_scores) if end_idx < start_idx: end_idx = start_idx # Get answer tokens input_ids = encoding['input_ids'][0] answer_tokens = input_ids[start_idx:end_idx+1] answer = self.processor.tokenizer.decode(answer_tokens, skip_special_tokens=True) # Calculate confidence confidence = (torch.softmax(start_scores, dim=-1).max() + torch.softmax(end_scores, dim=-1).max()) / 2 confidence = confidence.item() details = { "question": question, "start_position": start_idx.item(), "end_position": end_idx.item(), "total_words": len(words), "model_used": f"microsoft/layoutlmv3-{model_size}" } return (answer, confidence, details) # Node mappings NODE_CLASS_MAPPINGS = { "LayoutLMv3DocumentAnalysisNode": LayoutLMv3DocumentAnalysisNode, "LayoutLMv3QuestionAnsweringNode": LayoutLMv3QuestionAnsweringNode, } NODE_DISPLAY_NAME_MAPPINGS = { "LayoutLMv3DocumentAnalysisNode": "LayoutLMv3 Document Analysis", "LayoutLMv3QuestionAnsweringNode": "LayoutLMv3 Question Answering", }