EricRollei-PDF-Tools/oldfiles/enhanced_layout_parser_v02.py

import os

# Surya OCR performance optimizations for high-end hardware (24GB VRAM, 128GB RAM)
os.environ['RECOGNITION_BATCH_SIZE'] = '768'
os.environ['DETECTOR_BATCH_SIZE'] = '54'
os.environ['LAYOUT_BATCH_SIZE'] = '48'
os.environ['TABLE_REC_BATCH_SIZE'] = '96'
os.environ['TORCH_DEVICE'] = 'cuda'

import torch
import numpy as np
from PIL import Image, ImageDraw
from typing import List, Dict, Any, Tuple, Optional
import cv2

# ComfyUI imports
try:
    from comfy.utils import common_upscale
    from comfy.model_management import get_torch_device
    import folder_paths
    COMFYUI_BASE_PATH = os.path.dirname(folder_paths.models_dir)
    
    # Helper functions for tensor conversion
    def tensor_to_PIL(tensor):
        return Image.fromarray(np.clip(255. * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
    
    def PIL_to_tensor(pil_image):
        return torch.from_numpy(np.array(pil_image).astype(np.float32) / 255.0).unsqueeze(0)
    
except ImportError:
    print("ComfyUI imports not available, using fallback")
    COMFYUI_BASE_PATH = "."
    
    def tensor_to_PIL(tensor):
        if isinstance(tensor, torch.Tensor):
            return Image.fromarray(np.clip(255. * tensor.cpu().numpy().squeeze(), 0, 255).astype(np.uint8))
        return tensor
    
    def PIL_to_tensor(pil_image):
        return torch.from_numpy(np.array(pil_image).astype(np.float32) / 255.0).unsqueeze(0)

class EnhancedLayoutParserNode:
    """
    Multi-modal layout analysis combining:
    - LayoutLMv3 for semantic understanding
    - Florence2 for image detection  
    - Surya-OCR for precise text layout
    - CV2 for geometric analysis
    - Logic fusion for optimal results
    
    Optimized for high-end hardware (24GB VRAM, 128GB RAM)
    """
    
    @classmethod  
    def INPUT_TYPES(cls):
        return {
            "required": {
                "image": ("IMAGE",),
                "primary_ocr": (["surya", "paddleocr", "easyocr"], {"default": "surya"}),
                "enable_florence2": ("BOOLEAN", {"default": True}),
                "enable_layoutlmv3": ("BOOLEAN", {"default": True}),
                "enable_cv2_backup": ("BOOLEAN", {"default": False}),
                "extract_images": ("BOOLEAN", {"default": True}),
                "extract_text": ("BOOLEAN", {"default": True}),
                "fusion_strategy": (["conservative", "aggressive", "hybrid"], {"default": "aggressive"}),
            },
            "optional": {
                "florence2_image_prompt": ("STRING", {
                    "default": "rectangular images in page",
                    "multiline": True
                }),
                "florence2_text_prompt": ("STRING", {
                    "default": "text, caption, paragraph, title",
                    "multiline": True
                }),
                "min_image_area": ("INT", {"default": 10000, "min": 1000, "max": 100000}),
                "min_text_area": ("INT", {"default": 500, "min": 100, "max": 10000}),
                # Separate confidence controls for each method
                "surya_confidence": ("FLOAT", {"default": 0.3, "min": 0.1, "max": 1.0, "step": 0.05}),
                "florence2_confidence": ("FLOAT", {"default": 0.5, "min": 0.1, "max": 1.0, "step": 0.05}),
                "cv2_confidence": ("FLOAT", {"default": 0.7, "min": 0.1, "max": 1.0, "step": 0.05}),
                # Text grouping controls - ENHANCED
                "text_grouping_mode": (["individual_lines", "grouped_blocks", "smart_hybrid", "florence2_priority"], {"default": "florence2_priority"}),
                "text_merge_distance": ("INT", {"default": 25, "min": 5, "max": 100, "step": 5}),  # Made more granular
                # Add text recognition control
                "include_text_recognition": ("BOOLEAN", {"default": True}),
                # Enhanced debugging
                "enable_debug_logging": ("BOOLEAN", {"default": True}),
                # High-performance batch size controls
                "surya_recognition_batch": ("INT", {"default": 768, "min": 128, "max": 2048}),
                "surya_detector_batch": ("INT", {"default": 54, "min": 16, "max": 128}),
                # ADVANCED TEXT GROUPING CONTROLS
                "smart_grouping_enabled": ("BOOLEAN", {"default": True}),
                "vertical_merge_threshold": ("INT", {"default": 15, "min": 5, "max": 50}),  # For line spacing
                "horizontal_merge_threshold": ("INT", {"default": 10, "min": 5, "max": 30}),  # For word spacing
            }
        }

    RETURN_TYPES = ("IMAGE", "LIST", "LIST", "LIST", "LIST", "STRING", "DICT")
    RETURN_NAMES = ("overlay_image", "image_boxes", "text_boxes", "image_crops", "text_crops", "extracted_text", "analysis_results")

    FUNCTION = "analyze_layout"
    CATEGORY = "Enhanced Layout/Analysis"
    
    def __init__(self):
        self.surya_ocr = None
        self.florence2_detector = None
        self.layoutlmv3_model = None
        self.current_ocr = None
        print("🚀 Enhanced Layout Parser initialized with high-performance settings")
        print(f"   Recognition batch size: {os.environ.get('RECOGNITION_BATCH_SIZE', 'default')}")
        print(f"   Detector batch size: {os.environ.get('DETECTOR_BATCH_SIZE', 'default')}")
        print(f"   CUDA device: {os.environ.get('TORCH_DEVICE', 'auto')}")
        
    def _debug_surya_import(self):
        """Debug Surya import issues"""
        try:
            print("🔍 Debugging Surya import...")
            
            # Try basic import
            import surya
            print(f"✅ Basic surya import successful, version: {getattr(surya, '__version__', 'unknown')}")
            
            # Try specific imports
            from surya.ocr import run_ocr
            print("✅ surya.ocr.run_ocr import successful")
            
            from surya.model.detection.model import load_model as load_det_model
            print("✅ surya detection model import successful")
            
            from surya.model.recognition.model import load_model as load_rec_model  
            print("✅ surya recognition model import successful")
            
            return True
            
        except ImportError as e:
            print(f"❌ Surya import failed: {e}")
            print("💡 Try: pip install surya-ocr")
            return False
        except Exception as e:
            print(f"❌ Unexpected Surya error: {e}")
            return False

    def _init_surya_ocr(self, language="en"):
        """Initialize Surya-OCR with high-performance settings"""
        try:
            print("🔧 Initializing Surya-OCR with optimized batch sizes...")
            
            # Debug import first
            if not self._debug_surya_import():
                return False
            
            from surya.ocr import run_ocr
            from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor
            from surya.model.recognition.model import load_model as load_rec_model
            from surya.model.recognition.processor import load_processor as load_rec_processor
            
            # Load Surya models with CUDA optimization
            print("🔧 Loading detection models...")
            try:
                det_processor, det_model = load_det_processor(), load_det_model()
                print("✅ Detection models loaded successfully")
            except Exception as e:
                print(f"❌ Failed to load detection models: {e}")
                return False
            
            print("🔧 Loading recognition models...")
            try:
                rec_model, rec_processor = load_rec_model(), load_rec_processor()
                print("✅ Recognition models loaded successfully")
            except Exception as e:
                print(f"❌ Failed to load recognition models: {e}")
                return False
            
            # Move models to CUDA for maximum performance
            if torch.cuda.is_available():
                try:
                    det_model = det_model.cuda()
                    rec_model = rec_model.cuda()
                    print("✅ Models moved to CUDA")
                except Exception as e:
                    print(f"⚠️ Failed to move models to CUDA: {e}")
                    print("   Continuing with CPU...")
            
            self.surya_ocr = {
                'det_model': det_model,
                'det_processor': det_processor,
                'rec_model': rec_model,
                'rec_processor': rec_processor
            }
            
            print("✅ Surya-OCR initialized successfully with high-performance settings")
            if torch.cuda.is_available():
                print(f"   Available GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f}GB")
            return True
            
        except ImportError:
            print("❌ Surya-OCR not available. Install with: pip install surya-ocr")
            return False
        except Exception as e:
            print(f"❌ Error initializing Surya-OCR: {e}")
            import traceback
            traceback.print_exc()
            return False

    def _init_florence2(self):
        """Initialize Florence2 detector with high-performance settings"""
        try:
            from PDF_tools import Florence2RectangleDetector
            
            print("🔧 Initializing Florence2 with large model for maximum accuracy...")
            self.florence2_detector = Florence2RectangleDetector(
                model_name="CogFlorence-2.2-Large",  # Use large model with your VRAM
                comfyui_base_path=COMFYUI_BASE_PATH,
                min_box_area=1000
                # Remove device parameter - not supported by this class
            )
            print("✅ Florence2 detector initialized with CogFlorence-2.2-Large")
            return True
        except Exception as e:
            print(f"❌ Error initializing Florence2: {e}")
            # Try with base model as fallback
            try:
                print("🔧 Trying Florence2 with base model...")
                self.florence2_detector = Florence2RectangleDetector(
                    model_name="microsoft/Florence-2-base",
                    comfyui_base_path=COMFYUI_BASE_PATH,
                    min_box_area=1000
                )
                print("✅ Florence2 detector initialized with base model")
                return True
            except Exception as e2:
                print(f"❌ Florence2 base model also failed: {e2}")
                return False

    def _init_layoutlmv3(self):
        """Initialize LayoutLMv3 for semantic document understanding"""
        try:
            from transformers import LayoutLMv3Processor, LayoutLMv3ForTokenClassification
            
            print("🔧 Initializing LayoutLMv3 for semantic analysis...")
            
            # Load LayoutLMv3 model for document understanding
            self.layoutlmv3_processor = LayoutLMv3Processor.from_pretrained(
                "microsoft/layoutlmv3-base",
                apply_ocr=False  # We'll provide our own OCR results
            )
            
            self.layoutlmv3_model = LayoutLMv3ForTokenClassification.from_pretrained(
                "microsoft/layoutlmv3-base"
            )
            
            # Move to CUDA if available
            if torch.cuda.is_available():
                self.layoutlmv3_model = self.layoutlmv3_model.cuda()
                print("✅ LayoutLMv3 moved to CUDA")
            
            print("✅ LayoutLMv3 initialized successfully")
            return True
            
        except Exception as e:
            print(f"❌ LayoutLMv3 initialization failed: {e}")
            return False

    def _init_fallback_ocr(self, ocr_engine="paddleocr"):
        """Initialize fallback OCR when Surya is not available"""
        try:
            if ocr_engine == "paddleocr":
                from paddleocr import PaddleOCR
                # Remove show_log parameter - it's not supported in newer versions
                self.fallback_ocr = PaddleOCR(use_angle_cls=True, lang='en')
                print(f"✅ Fallback OCR (PaddleOCR) initialized")
                return True
            elif ocr_engine == "easyocr":
                import easyocr
                self.fallback_ocr = easyocr.Reader(['en'])
                print(f"✅ Fallback OCR (EasyOCR) initialized")
                return True
            else:
                return False
        except Exception as e:
            print(f"❌ Fallback OCR initialization failed: {e}")
            return False

    def _run_fallback_ocr(self, image: Image.Image, ocr_engine="paddleocr") -> List[Dict]:
        """Run fallback OCR when Surya is not available"""
        try:
            image_array = np.array(image)
            
            if ocr_engine == "paddleocr":
                results = self.fallback_ocr.ocr(image_array, cls=True)
                ocr_results = []
                if results and results[0]:  # PaddleOCR can return None
                    for line in results[0]:
                        bbox, (text, confidence) = line
                        x1 = int(min([point[0] for point in bbox]))
                        y1 = int(min([point[1] for point in bbox]))
                        x2 = int(max([point[0] for point in bbox]))
                        y2 = int(max([point[1] for point in bbox]))
                        
                        ocr_results.append({
                            'text': text,
                            'bbox': [x1, y1, x2, y2],
                            'confidence': confidence,
                            'source': 'paddleocr'
                        })
                        
            elif ocr_engine == "easyocr":
                results = self.fallback_ocr.readtext(image_array)
                ocr_results = []
                for (bbox, text, confidence) in results:
                    # EasyOCR returns bbox as [[x1,y1], [x2,y1], [x2,y2], [x1,y2]]
                    x1 = int(min([point[0] for point in bbox]))
                    y1 = int(min([point[1] for point in bbox]))
                    x2 = int(max([point[0] for point in bbox]))
                    y2 = int(max([point[1] for point in bbox]))
                    
                    ocr_results.append({
                        'text': text,
                        'bbox': [x1, y1, x2, y2],
                        'confidence': confidence,
                        'source': 'easyocr'
                    })
            
            print(f"✅ Fallback OCR ({ocr_engine}) processed {len(ocr_results)} text regions")
            return ocr_results
            
        except Exception as e:
            print(f"❌ Fallback OCR ({ocr_engine}) failed: {e}")
            return []

    def _run_surya_ocr(self, image: Image.Image) -> List[Dict]:
        """Run Surya-OCR with high-performance batch processing"""
        try:
            from surya.ocr import run_ocr
            
            # For large images, process in parallel batches
            image_width, image_height = image.size
            
            # Use larger batch processing for your hardware
            if image_width * image_height > 2000000:  # Large image
                print(f"🚀 Processing large image ({image_width}x{image_height}) with optimized batching")
            
            # Run OCR with optimized settings
            predictions = run_ocr(
                [image], 
                [["en"]], 
                self.surya_ocr['det_model'],
                self.surya_ocr['det_processor'], 
                self.surya_ocr['rec_model'],
                self.surya_ocr['rec_processor']
            )
            
            # Convert to our format
            ocr_results = []
            for prediction in predictions:
                for text_line in prediction.text_lines:
                    bbox = text_line.bbox
                    text = text_line.text
                    confidence = text_line.confidence if hasattr(text_line, 'confidence') else 0.9
                    
                    ocr_results.append({
                        'text': text,
                        'bbox': [int(bbox[0]), int(bbox[1]), int(bbox[2]), int(bbox[3])],
                        'confidence': confidence,
                        'source': 'surya'
                    })
            
            print(f"✅ Surya-OCR processed {len(ocr_results)} text regions")
            return ocr_results
            
        except Exception as e:
            print(f"❌ Surya-OCR failed: {e}")
            return []
    
    def _run_florence2_detection(self, image: Image.Image, 
                               image_prompt: str = "rectangular images in page",
                               text_prompt: str = "text, caption, paragraph, title") -> Tuple[List[Dict], List[Dict]]:
        """Run Florence2 with optimized prompts for both image and text detection"""
        try:
            # Detect images with optimized prompt
            image_boxes, _ = self.florence2_detector.detect_rectangles(
                image=image,
                text_input=image_prompt,
                return_mask=False,
                keep_model_loaded=True
            )
            
            # Convert to our format with integer coordinates
            florence2_images = []
            for box in image_boxes:
                bbox = [int(box.x1), int(box.y1), int(box.x2), int(box.y2)]
                florence2_images.append({
                    'bbox': bbox,
                    'label': box.label,
                    'confidence': box.confidence if box.confidence else 0.8,
                    'area': (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]),
                    'source': 'florence2'
                })
            
            # Detect text regions with optimized prompt
            text_boxes, _ = self.florence2_detector.detect_rectangles(
                image=image,
                text_input=text_prompt,
                return_mask=False,
                keep_model_loaded=True
            )
            
            florence2_text = []
            for box in text_boxes:
                bbox = [int(box.x1), int(box.y1), int(box.x2), int(box.y2)]
                florence2_text.append({
                    'bbox': bbox,
                    'label': box.label,
                    'confidence': box.confidence if box.confidence else 0.8,
                    'area': (bbox[2] - bbox[0]) * (bbox[3] - bbox[1]),
                    'source': 'florence2'
                })
            
            print(f"✅ Florence2 found {len(florence2_images)} images, {len(florence2_text)} text regions")
            return florence2_images, florence2_text
            
        except Exception as e:
            print(f"❌ Florence2 detection failed: {e}")
            return [], []
    
    def _run_cv2_analysis(self, image: Image.Image) -> Tuple[List[Dict], List[Dict]]:
        """CV2-based geometric analysis as backup"""
        try:
            img_array = np.array(image)
            gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
            
            # Image detection using contours
            edges = cv2.Canny(gray, 50, 150)
            contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            cv2_images = []
            for contour in contours:
                area = cv2.contourArea(contour)
                if area < 5000:  # Filter small contours
                    continue
                
                x, y, w, h = cv2.boundingRect(contour)
                aspect_ratio = w / h
                
                # Filter for image-like regions (rectangular, reasonable aspect ratio)
                if 0.3 < aspect_ratio < 3.0 and w > 100 and h > 100:
                    cv2_images.append({
                        'bbox': [x, y, x + w, y + h],
                        'label': 'image_region',
                        'confidence': 0.6,
                        'area': area,
                        'source': 'cv2'
                    })
            
            # Text detection using MSER (more sensitive to text regions)
            mser = cv2.MSER_create()
            regions, _ = mser.detectRegions(gray)
            
            cv2_text = []
            for region in regions:
                if len(region) < 50:  # Filter very small regions
                    continue
                
                hull = cv2.convexHull(region)
                x, y, w, h = cv2.boundingRect(hull)
                
                # Filter for text-like regions
                if w > 20 and h > 10 and w/h > 2:  # Text tends to be wider than tall
                    cv2_text.append({
                        'bbox': [x, y, x + w, y + h],
                        'label': 'text_region',
                        'confidence': 0.5,
                        'area': w * h,
                        'source': 'cv2'
                    })
            
            return cv2_images, cv2_text
            
        except Exception as e:
            print(f"❌ CV2 analysis failed: {e}")
            return [], []
    def _run_layoutlmv3_analysis(self, image: Image.Image, text_boxes: List[Dict]) -> Dict[str, List[Dict]]:
        """Use LayoutLMv3 to classify and confirm text/image regions semantically"""
        try:
            if self.layoutlmv3_model is None:
                if not self._init_layoutlmv3():
                    return {"confirmed_text": [], "confirmed_images": [], "semantic_labels": []}
            
            # Prepare OCR data for LayoutLMv3
            words = []
            boxes = []
            
            for text_box in text_boxes:
                if text_box.get('text', '').strip():
                    # Split text into words and assign same bbox to all words
                    text_words = text_box['text'].split()
                    for word in text_words:
                        words.append(word)
                        boxes.append(text_box['bbox'])  # [x1, y1, x2, y2]
            
            if not words:
                return {"confirmed_text": [], "confirmed_images": [], "semantic_labels": []}
            
            # Process with LayoutLMv3
            encoding = self.layoutlmv3_processor(
                image, 
                words, 
                boxes=boxes, 
                return_tensors="pt",
                padding="max_length",
                truncation=True
            )
            
            # Move to CUDA if available
            if torch.cuda.is_available():
                encoding = {k: v.cuda() for k, v in encoding.items()}
            
            # Run inference
            with torch.no_grad():
                outputs = self.layoutlmv3_model(**encoding)
                predictions = outputs.logits.argmax(-1).squeeze().tolist()
            
            # Map predictions to semantic labels
            # LayoutLMv3 typically outputs: 0=Other, 1=Header, 2=Question, 3=Answer, etc.
            label_map = {
                0: "other",
                1: "header", 
                2: "title",
                3: "text",
                4: "list",
                5: "table",
                6: "figure"
            }
            
            # Group results by semantic meaning
            confirmed_text = []
            confirmed_images = []
            semantic_labels = []
            
            current_idx = 0
            for text_box in text_boxes:
                if text_box.get('text', '').strip():
                    text_words = text_box['text'].split()
                    
                    # Get most common prediction for this text box
                    box_predictions = predictions[current_idx:current_idx + len(text_words)]
                    most_common_pred = max(set(box_predictions), key=box_predictions.count)
                    semantic_label = label_map.get(most_common_pred, "other")
                    
                    # Enhance the text box with semantic information
                    enhanced_box = {
                        **text_box,
                        'semantic_label': semantic_label,
                        'semantic_confidence': box_predictions.count(most_common_pred) / len(box_predictions)
                    }
                    
                    # Classify as text or potential figure based on semantic label
                    if semantic_label in ["figure", "table"]:
                        confirmed_images.append(enhanced_box)
                    else:
                        confirmed_text.append(enhanced_box)
                    
                    semantic_labels.append({
                        'bbox': text_box['bbox'],
                        'label': semantic_label,
                        'confidence': enhanced_box['semantic_confidence']
                    })
                    
                    current_idx += len(text_words)
            
            print(f"✅ LayoutLMv3 analyzed {len(semantic_labels)} regions")
            return {
                "confirmed_text": confirmed_text,
                "confirmed_images": confirmed_images, 
                "semantic_labels": semantic_labels
            }
            
        except Exception as e:
            print(f"❌ LayoutLMv3 analysis failed: {e}")
            return {"confirmed_text": [], "confirmed_images": [], "semantic_labels": []}

    def _smart_fusion_with_layoutlmv3(self, 
                                     ocr_results: List[Dict],
                                     florence2_images: List[Dict],
                                     florence2_text: List[Dict], 
                                     cv2_images: List[Dict],
                                     cv2_text: List[Dict],
                                     strategy: str,
                                     min_image_area: int,
                                     min_text_area: int,
                                     surya_confidence: float = 0.3,
                                     florence2_confidence: float = 0.5,
                                     cv2_confidence: float = 0.7,
                                     text_grouping_mode: str = "smart_hybrid",
                                     enable_layoutlmv3: bool = True) -> Tuple[List[Dict], List[Dict]]:
        """Enhanced fusion with LayoutLMv3 validation of Florence2 grouping"""
        
        print(f"🧠 === FUSION WITH LAYOUTLMV3 DEBUG ===")
        print(f"🧠 Input counts: OCR={len(ocr_results)}, F2_img={len(florence2_images)}, F2_txt={len(florence2_text)}")
        
        # First, do the basic fusion with enhanced Florence2 filtering
        final_images, final_text = self._smart_fusion(
            ocr_results, florence2_images, florence2_text, cv2_images, cv2_text,
            strategy, min_image_area, min_text_area,
            surya_confidence, florence2_confidence, cv2_confidence, text_grouping_mode
        )
        
        print(f"🧠 After basic fusion: Images={len(final_images)}, Text={len(final_text)}")
        
        # Then use LayoutLMv3 to validate and enhance Florence2's grouping decisions
        if enable_layoutlmv3 and final_text:
            print("🧠 Running LayoutLMv3 to validate Florence2 text grouping...")
            
            # Separate Florence2 text boxes for special validation
            florence2_text_boxes = [txt for txt in final_text if 'florence2' in txt['detection_methods']]
            other_text_boxes = [txt for txt in final_text if 'florence2' not in txt['detection_methods']]
            
            print(f"🧠 LayoutLMv3 validation targets: {len(florence2_text_boxes)} Florence2 boxes, "
                  f"{len(other_text_boxes)} other boxes")
            
            # Filter text boxes that have actual text content for LayoutLMv3
            text_with_content = [txt for txt in final_text if txt.get('text', '').strip()]
            
            if text_with_content:
                layoutlmv3_results = self._run_layoutlmv3_analysis(self.current_image, text_with_content)
                
                semantic_confirmed_text = layoutlmv3_results["confirmed_text"]
                semantic_confirmed_images = layoutlmv3_results["confirmed_images"]
                semantic_labels = layoutlmv3_results["semantic_labels"]
                
                print(f"🧠 LayoutLMv3 results:")
                print(f"   - confirmed_text: {len(semantic_confirmed_text)}")
                print(f"   - confirmed_images: {len(semantic_confirmed_images)}")
                print(f"   - semantic_labels: {len(semantic_labels)}")
                
                # SPECIAL VALIDATION FOR FLORENCE2 TEXT GROUPING
                florence2_validated = 0
                florence2_rejected = 0
                
                for txt in final_text:
                    if 'florence2' in txt['detection_methods']:
                        # Find corresponding semantic result
                        semantic_validation = None
                        for sem_txt in semantic_confirmed_text:
                            overlap_ratio = self._calculate_overlap_ratio(txt['bbox'], sem_txt['bbox'])
                            if overlap_ratio > 0.3:  # Reasonable overlap
                                semantic_validation = sem_txt
                                break
                        
                        if semantic_validation:
                            semantic_label = semantic_validation.get('semantic_label', 'text')
                            semantic_confidence = semantic_validation.get('semantic_confidence', 0.5)
                            
                            # Validate Florence2's grouping decision
                            if semantic_label in ['text', 'header', 'title', 'list'] and semantic_confidence > 0.6:
                                # LayoutLMv3 confirms this is good text grouping
                                txt['semantic_label'] = semantic_label
                                txt['semantic_confidence'] = semantic_confidence
                                txt['layoutlmv3_validation'] = 'confirmed'
                                
                                # Boost confidence for validated Florence2 grouping
                                old_confidence = txt['confidence']
                                txt['confidence'] = min(txt['confidence'] * 1.3, 1.0)  # Bigger boost
                                florence2_validated += 1
                                
                                if 'layoutlmv3' not in txt['detection_methods']:
                                    txt['detection_methods'].append('layoutlmv3')
                                    txt['confidence_sources'].append(semantic_confidence)
                                
                                print(f"🧠 ✅ Validated Florence2 grouping: {semantic_label} "
                                      f"conf {old_confidence:.2f}→{txt['confidence']:.2f}")
                            
                            elif semantic_label in ['figure', 'table']:
                                # Florence2 grouped text that's actually an image region
                                txt['layoutlmv3_validation'] = 'reclassified_as_image'
                                txt['semantic_label'] = semantic_label
                                
                                # Move to images if it meets size criteria
                                if txt.get('area', 0) >= min_image_area:
                                    final_images.append({
                                        **txt,
                                        'detection_methods': ['florence2', 'layoutlmv3'],
                                        'confidence_sources': [txt['confidence'], semantic_confidence],
                                        'source': 'florence2_reclassified',
                                        'reclassification_reason': f'LayoutLMv3 identified as {semantic_label}'
                                    })
                                    print(f"🧠 🔄 Reclassified Florence2 text as {semantic_label}")
                                
                                # Mark for removal from text
                                txt['layoutlmv3_validation'] = 'remove_from_text'
                                florence2_rejected += 1
                            
                            else:
                                # Low confidence or unknown label
                                txt['layoutlmv3_validation'] = 'uncertain'
                                txt['semantic_label'] = semantic_label
                                txt['semantic_confidence'] = semantic_confidence
                        
                        else:
                            # No semantic validation found - might be problematic Florence2 grouping
                            txt['layoutlmv3_validation'] = 'no_semantic_match'
                            # Reduce confidence for unvalidated Florence2 text
                            txt['confidence'] = txt['confidence'] * 0.8
                
                # Remove Florence2 text boxes that were reclassified
                final_text = [txt for txt in final_text if txt.get('layoutlmv3_validation') != 'remove_from_text']
                
                # Process other text boxes normally
                text_enhanced_count = 0
                text_confidence_boosted = 0
                
                for txt in final_text:
                    if 'florence2' not in txt['detection_methods']:  # Non-Florence2 text
                        for sem_txt in semantic_confirmed_text:
                            overlap_ratio = self._calculate_overlap_ratio(txt['bbox'], sem_txt['bbox'])
                            if overlap_ratio > 0.5:
                                txt['semantic_label'] = sem_txt.get('semantic_label', 'text')
                                txt['semantic_confidence'] = sem_txt.get('semantic_confidence', 0.5)
                                text_enhanced_count += 1
                                
                                if txt['semantic_confidence'] > 0.7:
                                    old_confidence = txt['confidence']
                                    txt['confidence'] = min(txt['confidence'] * 1.2, 1.0)
                                    text_confidence_boosted += 1
                                    
                                    if 'layoutlmv3' not in txt['detection_methods']:
                                        txt['detection_methods'].append('layoutlmv3')
                                        txt['confidence_sources'].append(txt['semantic_confidence'])
                                break
                
                # Add semantically detected images
                images_added_by_semantic = 0
                for sem_img in semantic_confirmed_images:
                    if (sem_img['semantic_label'] in ['figure', 'table'] and 
                        sem_img.get('area', 0) >= min_image_area):
                        
                        has_overlap = False
                        for existing_img in final_images:
                            overlap_ratio = self._calculate_overlap_ratio(sem_img['bbox'], existing_img['bbox'])
                            if overlap_ratio > 0.3:
                                has_overlap = True
                                break
                        
                        if not has_overlap:
                            final_images.append({
                                **sem_img,
                                'detection_methods': ['layoutlmv3'],
                                'confidence_sources': [sem_img.get('semantic_confidence', 0.8)],
                                'source': 'layoutlmv3_semantic'
                            })
                            images_added_by_semantic += 1
                
                print(f"🧠 LayoutLMv3 FLORENCE2 VALIDATION SUMMARY:")
                print(f"   - Florence2 text boxes validated: {florence2_validated}")
                print(f"   - Florence2 text boxes rejected/reclassified: {florence2_rejected}")
                print(f"   - Other text boxes enhanced: {text_enhanced_count}")
                print(f"   - Confidence boosts applied: {text_confidence_boosted}")
                print(f"   - Images added by semantic detection: {images_added_by_semantic}")
                
            else:
                print("🧠 No text content available for LayoutLMv3 analysis")
        
        print(f"🧠 Final result: Images={len(final_images)}, Text={len(final_text)}")
        print(f"🧠 === END LAYOUTLMV3 FUSION DEBUG ===")
        return final_images, final_text

    def _smart_fusion(self, 
                     ocr_results: List[Dict],
                     florence2_images: List[Dict],
                     florence2_text: List[Dict], 
                     cv2_images: List[Dict],
                     cv2_text: List[Dict],
                     strategy: str,
                     min_image_area: int,
                     min_text_area: int,
                     surya_confidence: float = 0.3,
                     florence2_confidence: float = 0.5,
                     cv2_confidence: float = 0.7,
                     text_grouping_mode: str = "smart_hybrid") -> Tuple[List[Dict], List[Dict]]:
        """Smart fusion with enhanced Florence2 text filtering"""
        
        print(f"🔍 === SMART FUSION DEBUG ===")
        print(f"🔍 Fusion input: OCR={len(ocr_results)}, F2_img={len(florence2_images)}, "
              f"F2_txt={len(florence2_text)}, CV2_img={len(cv2_images)}, CV2_txt={len(cv2_text)}")
        
        # Calculate image dimensions for filtering
        image_width = getattr(self.current_image, 'width', 1000)
        image_height = getattr(self.current_image, 'height', 1000)
        total_image_area = image_width * image_height
        
        print(f"🔍 Image dimensions: {image_width}x{image_height}, total area: {total_image_area}")
        
        # 1. PRIORITIZE FLORENCE2 IMAGES (best performer)
        final_images = []
        florence2_images_added = 0
        for f2_img in florence2_images:
            if (f2_img['area'] >= min_image_area and 
                f2_img['confidence'] >= florence2_confidence):
                final_images.append({
                    **f2_img,
                    'detection_methods': ['florence2'],
                    'confidence_sources': [f2_img['confidence']]
                })
                florence2_images_added += 1
        
        print(f"🔍 Florence2 images added: {florence2_images_added}/{len(florence2_images)}")
        
        # 2. PROCESS TEXT BASED ON GROUPING MODE with enhanced Florence2 filtering
        final_text = []
        surya_text_added = 0
        florence2_text_added = 0
        florence2_text_filtered = 0
        
        # Get grouping parameters
        merge_distance = getattr(self, 'text_merge_distance', 25)
        smart_grouping = getattr(self, 'smart_grouping_enabled', True)
        v_threshold = getattr(self, 'vertical_merge_threshold', 15)
        h_threshold = getattr(self, 'horizontal_merge_threshold', 10)
        
        if text_grouping_mode == "florence2_priority":
            print(f"🔍 Florence2 Priority mode with smart filtering...")
            
            # ENHANCED FLORENCE2 TEXT FILTERING
            filtered_florence2_text = []
            for f2_txt in florence2_text:
                # Calculate relative size to image
                relative_area = f2_txt['area'] / total_image_area
                
                # Filter criteria
                too_large = relative_area > 0.8  # More than 80% of image
                too_small = f2_txt['area'] < min_text_area
                low_confidence = f2_txt['confidence'] < florence2_confidence
                
                # Aspect ratio check (whole page boxes tend to be very tall or very wide)
                bbox = f2_txt['bbox']
                width = bbox[2] - bbox[0]
                height = bbox[3] - bbox[1]
                aspect_ratio = width / height if height > 0 else 1
                unusual_aspect = aspect_ratio > 5 or aspect_ratio < 0.1  # Very wide or very tall
                
                # Position check (whole page boxes often start near 0,0)
                near_origin = bbox[0] < 50 and bbox[1] < 50
                covers_most_width = width > image_width * 0.9
                covers_most_height = height > image_height * 0.9
                
                if too_large or too_small or low_confidence:
                    florence2_text_filtered += 1
                    print(f"🔍 Filtered Florence2 text: area={f2_txt['area']} ({relative_area:.1%}), "
                          f"conf={f2_txt['confidence']:.2f}, reasons: "
                          f"{'too_large ' if too_large else ''}"
                          f"{'too_small ' if too_small else ''}"
                          f"{'low_conf ' if low_confidence else ''}")
                    continue
                
                if unusual_aspect and (near_origin or covers_most_width or covers_most_height):
                    florence2_text_filtered += 1
                    print(f"🔍 Filtered Florence2 text (whole-page): aspect={aspect_ratio:.2f}, "
                          f"pos=({bbox[0]},{bbox[1]}), size=({width}x{height})")
                    continue
                
                # If it passes all filters, it's good Florence2 text
                filtered_florence2_text.append(f2_txt)
            
            print(f"🔍 Florence2 text filtering: {len(florence2_text)} → {len(filtered_florence2_text)} "
                  f"(filtered out {florence2_text_filtered})")
            
            # Add filtered Florence2 text first
            for f2_txt in filtered_florence2_text:
                final_text.append({
                    **f2_txt,
                    'detection_methods': ['florence2'],
                    'confidence_sources': [f2_txt['confidence']],
                    'text': '',  # Florence2 doesn't provide actual text content
                    'text_extraction_method': 'florence2_bbox_only'
                })
                florence2_text_added += 1
            
            # Then supplement with Surya grouped text for actual content and gaps
            grouped_ocr = self._group_nearby_text(
                ocr_results, 
                merge_distance=merge_distance,
                smart_grouping=smart_grouping,
                vertical_threshold=v_threshold,
                horizontal_threshold=h_threshold
            )
            
            for group in grouped_ocr:
                if (group.get('area', 0) >= min_text_area // 2 and 
                    group['confidence'] >= surya_confidence):
                    
                    # Check if this overlaps significantly with Florence2 text
                    best_florence_match = None
                    best_overlap = 0
                    
                    for f2_txt in final_text:
                        if 'florence2' in f2_txt['detection_methods']:
                            overlap_ratio = self._calculate_overlap_ratio(group['bbox'], f2_txt['bbox'])
                            if overlap_ratio > best_overlap:
                                best_overlap = overlap_ratio
                                best_florence_match = f2_txt
                    
                    if best_overlap > 0.3:  # Significant overlap with Florence2
                        # Enhance Florence2 detection with Surya text content
                        if best_florence_match and not best_florence_match.get('text'):
                            best_florence_match['text'] = group['text']
                            best_florence_match['detection_methods'].append('surya_grouped')
                            best_florence_match['confidence_sources'].append(group['confidence'])
                            best_florence_match['text_extraction_method'] = 'florence2_bbox_surya_text'
                            print(f"🔍 Enhanced Florence2 with Surya text: '{group['text'][:50]}...'")
                    else:
                        # Add as separate Surya text (fills gaps)
                        final_text.append({
                            **group,
                            'detection_methods': ['surya_grouped'],
                            'confidence_sources': [group['confidence']],
                            'text_extraction_method': 'surya_only'
                        })
                        surya_text_added += 1
        
        elif text_grouping_mode == "individual_lines":
            # Individual Surya detections for precise boundaries
            for ocr_result in ocr_results:
                if (ocr_result.get('area', 0) >= min_text_area // 2 and 
                    ocr_result['confidence'] >= surya_confidence):
                    final_text.append({
                        **ocr_result,
                        'detection_methods': [ocr_result['source']],
                        'confidence_sources': [ocr_result['confidence']],
                        'text_extraction_method': 'surya_individual'
                    })
                    surya_text_added += 1
        
        elif text_grouping_mode == "grouped_blocks":
            # Surya grouped blocks only
            grouped_ocr = self._group_nearby_text(
                ocr_results, 
                merge_distance=merge_distance,
                smart_grouping=smart_grouping,
                vertical_threshold=v_threshold,
                horizontal_threshold=h_threshold
            )
            for group in grouped_ocr:
                if (group.get('area', 0) >= min_text_area and 
                    group['confidence'] >= surya_confidence):
                    final_text.append({
                        **group,
                        'detection_methods': ['surya_grouped'],
                        'confidence_sources': [group['confidence']],
                        'text_extraction_method': 'surya_grouped'
                    })
                    surya_text_added += 1
        
        else:  # smart_hybrid
            # Use strategy-based approach
            if strategy == "conservative":
                # High confidence individual lines
                for ocr_result in ocr_results:
                    if (ocr_result.get('area', 0) >= min_text_area and 
                        ocr_result['confidence'] >= surya_confidence + 0.2):
                        final_text.append({
                            **ocr_result,
                            'detection_methods': [ocr_result['source']],
                            'confidence_sources': [ocr_result['confidence']],
                            'text_extraction_method': 'surya_conservative'
                        })
                        surya_text_added += 1
            else:
                # Grouped blocks for aggressive/hybrid
                grouped_ocr = self._group_nearby_text(
                    ocr_results, 
                    merge_distance=merge_distance,
                    smart_grouping=smart_grouping,
                    vertical_threshold=v_threshold,
                    horizontal_threshold=h_threshold
                )
                for group in grouped_ocr:
                    if (group.get('area', 0) >= min_text_area // 2 and 
                        group['confidence'] >= surya_confidence):
                        final_text.append({
                            **group,
                            'detection_methods': ['surya_grouped'],
                            'confidence_sources': [group['confidence']],
                            'text_extraction_method': 'surya_hybrid'
                        })
                        surya_text_added += 1
        
        print(f"🔍 Text results: Florence2={florence2_text_added}, Surya={surya_text_added}")
        print(f"🔍 Florence2 filtering removed {florence2_text_filtered} problematic boxes")
        
        # 3. APPLY STRATEGY-SPECIFIC FILTERING
        pre_filter_images = len(final_images)
        pre_filter_text = len(final_text)
        
        if strategy == "conservative":
            final_images = [img for img in final_images if 
                          len(img['detection_methods']) > 1 or img['confidence'] > 0.8]
            final_text = [txt for txt in final_text if 
                         len(txt['detection_methods']) > 1 or txt['confidence'] > 0.8]
        elif strategy == "aggressive":
            pass  # Keep most detections
        else:  # hybrid
            for img in final_images:
                if len(img['detection_methods']) > 1:
                    img['confidence'] = max(img['confidence_sources'])
                else:
                    img['confidence'] = img['confidence_sources'][0] * 0.9
            
            for txt in final_text:
                if len(txt['detection_methods']) > 1:
                    # Florence2 + Surya combo gets boost
                    txt['confidence'] = max(txt['confidence_sources'])
                else:
                    if 'surya' in txt['detection_methods'][0]:
                        txt['confidence'] = txt['confidence_sources'][0] * 0.95
                    else:
                        txt['confidence'] = txt['confidence_sources'][0] * 0.8
        
        # 4. FINAL FILTERING
        final_images = [img for img in final_images if img['confidence'] > 0.4]
        final_text = [txt for txt in final_text if txt['confidence'] > 0.3]
        
        print(f"🔍 Strategy filtering: Images {pre_filter_images}→{len(final_images)}, Text {pre_filter_text}→{len(final_text)}")
        print(f"🎯 Final fusion result: Images={len(final_images)}, Text={len(final_text)}")
        print(f"🔍 === END SMART FUSION DEBUG ===")
        
        return final_images, final_text
    
    def _group_nearby_text(self, ocr_results: List[Dict], merge_distance: int = 25, smart_grouping: bool = True, vertical_threshold: int = 15, horizontal_threshold: int = 10) -> List[Dict]:
        """Enhanced text grouping with smarter distance calculations"""
        if not ocr_results:
            return []
        
        print(f"🔍 Grouping {len(ocr_results)} OCR results with merge_distance={merge_distance}")
        
        groups = []
        used = set()
        
        for i, result in enumerate(ocr_results):
            if i in used:
                continue
            
            # Start new group
            group_texts = [result['text']]
            group_bbox = result['bbox'][:]
            group_confidence = result['confidence']
            group_indices = {i}
            group_areas = [result.get('area', 0)]
            
            # Find nearby text with smart grouping
            for j, other in enumerate(ocr_results):
                if j <= i or j in used:
                    continue
                
                if smart_grouping:
                    # Use different thresholds for vertical vs horizontal merging
                    should_merge = self._smart_should_merge(result['bbox'], other['bbox'], 
                                                          vertical_threshold, horizontal_threshold, merge_distance)
                else:
                    # Use simple distance
                    distance = self._calculate_bbox_distance(result['bbox'], other['bbox'])
                    should_merge = distance <= merge_distance
                
                if should_merge:
                    group_texts.append(other['text'])
                    group_bbox = self._merge_bboxes(group_bbox, other['bbox'])
                    group_confidence = min(group_confidence, other['confidence'])
                    group_indices.add(j)
                    group_areas.append(other.get('area', 0))
            
            # Add indices to used set
            used.update(group_indices)
            
            # Create group with enhanced metadata
            total_area = sum(group_areas)
            groups.append({
                'text': ' '.join(group_texts),
                'bbox': group_bbox,
                'confidence': group_confidence,
                'area': total_area,
                'source': 'surya_grouped',
                'component_count': len(group_texts),
                'merge_distance_used': merge_distance,
                'grouping_method': 'smart' if smart_grouping else 'simple'
            })
        
        print(f"✅ Grouped into {len(groups)} text blocks (reduction: {len(ocr_results)} → {len(groups)})")
        
        # Debug top groups
        for i, group in enumerate(groups[:3]):
            print(f"   Group {i}: {group['component_count']} components, area={group['area']}, text='{group['text'][:50]}...'")
        
        return groups

    def _smart_should_merge(self, bbox1: List[int], bbox2: List[int], 
                           vertical_threshold: int, horizontal_threshold: int, 
                           max_distance: int) -> bool:
        """Smart merging logic that considers text layout patterns"""
        x1_min, y1_min, x1_max, y1_max = bbox1
        x2_min, y2_min, x2_max, y2_max = bbox2
        
        # Calculate centers and dimensions
        x1_center = (x1_min + x1_max) / 2
        y1_center = (y1_min + y1_max) / 2
        x2_center = (x2_min + x2_max) / 2
        y2_center = (y2_min + y2_max) / 2
        
        h1 = y1_max - y1_min
        h2 = y2_max - y2_min
        w1 = x1_max - x1_min
        w2 = x2_max - x2_min
        
        # Calculate distances
        vertical_distance = abs(y1_center - y2_center)
        horizontal_distance = abs(x1_center - x2_center)
        
        # Check for vertical alignment (same column of text)
        horizontal_overlap = max(0, min(x1_max, x2_max) - max(x1_min, x2_min))
        vertical_overlap = max(0, min(y1_max, y2_max) - max(y1_min, y2_min))
        
        avg_width = (w1 + w2) / 2
        avg_height = (h1 + h2) / 2
        
        # Case 1: Vertically aligned text (same column)
        if horizontal_overlap > avg_width * 0.3:  # Significant horizontal overlap
            return vertical_distance <= vertical_threshold + avg_height * 0.5
        
        # Case 2: Horizontally aligned text (same line)  
        if vertical_overlap > avg_height * 0.3:  # Significant vertical overlap
            return horizontal_distance <= horizontal_threshold + avg_width * 0.2
        
        # Case 3: General proximity
        total_distance = (vertical_distance ** 2 + horizontal_distance ** 2) ** 0.5
        return total_distance <= max_distance


    def _calculate_overlap_ratio(self, bbox1: List[int], bbox2: List[int]) -> float:
        """Calculate overlap ratio between two bounding boxes"""
        x1_min, y1_min, x1_max, y1_max = bbox1
        x2_min, y2_min, x2_max, y2_max = bbox2
        
        # Calculate intersection
        inter_x_min = max(x1_min, x2_min)
        inter_y_min = max(y1_min, y2_min)
        inter_x_max = min(x1_max, x2_max)
        inter_y_max = min(y1_max, y2_max)
        
        if inter_x_max <= inter_x_min or inter_y_max <= inter_y_min:
            return 0.0
        
        inter_area = (inter_x_max - inter_x_min) * (inter_y_max - inter_y_min)
        bbox1_area = (x1_max - x1_min) * (y1_max - y1_min)
        bbox2_area = (x2_max - x2_min) * (y2_max - y2_min)
        
        union_area = bbox1_area + bbox2_area - inter_area
        return inter_area / union_area if union_area > 0 else 0.0
    
    def _calculate_bbox_distance(self, bbox1: List[int], bbox2: List[int]) -> float:
        """Calculate minimum distance between two bounding boxes"""
        x1_min, y1_min, x1_max, y1_max = bbox1
        x2_min, y2_min, x2_max, y2_max = bbox2
        
        # Calculate center points
        x1_center = (x1_min + x1_max) / 2
        y1_center = (y1_min + y1_max) / 2
        x2_center = (x2_min + x2_max) / 2
        y2_center = (y2_min + y2_max) / 2
        
        return ((x1_center - x2_center) ** 2 + (y1_center - y2_center) ** 2) ** 0.5
    
    def _merge_bboxes(self, bbox1: List[int], bbox2: List[int]) -> List[int]:
        """Merge two bounding boxes into one encompassing box"""
        return [
            min(bbox1[0], bbox2[0]),  # x_min
            min(bbox1[1], bbox2[1]),  # y_min
            max(bbox1[2], bbox2[2]),  # x_max
            max(bbox1[3], bbox2[3])   # y_max
        ]
    
    def _integrate_basic_surya_detection(self, image: Image.Image, confidence_threshold: float = 0.3) -> List[Dict]:
        """Use DetectionPredictor directly for better text detection coverage"""
        try:
            # Import the basic surya node
            from .basic_surya import SuryaOCRNode
            
            # Initialize if needed
            if not hasattr(self, 'basic_surya_node'):
                self.basic_surya_node = SuryaOCRNode()
            
            # Convert PIL to tensor format expected by basic_surya
            img_array = np.array(image).astype(np.float32) / 255.0
            img_tensor = torch.from_numpy(img_array).unsqueeze(0)
            
            # Run DETECTION ONLY for better coverage
            text_output, json_output, annotated_image, detection_data = self.basic_surya_node.process_ocr(
                img_tensor, 
                task_mode="detection_only",  # This is the key change!
                confidence_threshold=confidence_threshold
            )
            
            print(f"🔍 Debug - Surya Detection returned: text_output={bool(text_output)}, json_output={bool(json_output)}")
            print(f"🔍 Debug - Detection data type: {type(detection_data)}")
            
            # Convert detection_data to our format
            ocr_results = []
            
            # Method 1: Try detection_data directly
            if detection_data and isinstance(detection_data, list):
                print(f"🔍 Debug - Processing {len(detection_data)} detection items")
                for i, detection in enumerate(detection_data):
                    print(f"🔍 Debug - Detection {i}: keys={detection.keys() if isinstance(detection, dict) else 'not dict'}")
                    bbox = detection.get('bbox')
                    confidence = detection.get('confidence', 0.9)
                    
                    if bbox and confidence >= confidence_threshold:
                        area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                        ocr_results.append({
                            'text': '',  # Detection only, no text recognition
                            'bbox': bbox,
                            'confidence': confidence,
                            'area': area,
                            'source': 'surya_detection'
                        })
            
            # Method 2: Try JSON parsing
            elif json_output:
                print(f"🔍 Debug - Trying JSON parsing...")
                import json
                try:
                    results_dict = json.loads(json_output)
                    print(f"🔍 Debug - JSON structure: {list(results_dict.keys())}")
                    
                    detections = results_dict.get('detections', [])
                    if detections:
                        print(f"🔍 Debug - Found {len(detections)} detection pages")
                        for page_idx, page_detections in enumerate(detections):
                            if isinstance(page_detections, dict):
                                bboxes = page_detections.get('bboxes', [])
                                print(f"🔍 Debug - Page {page_idx} has {len(bboxes)} bboxes")
                                for bbox_idx, bbox_info in enumerate(bboxes):
                                    bbox = bbox_info.get('bbox')
                                    confidence = bbox_info.get('confidence', 0.9)
                                    
                                    if bbox and confidence >= confidence_threshold:
                                        area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                                        ocr_results.append({
                                            'text': '',
                                            'bbox': bbox,
                                            'confidence': confidence,
                                            'area': area,
                                            'source': 'surya_detection'
                                        })
                                        if bbox_idx < 3:  # Debug first few
                                            print(f"🔍 Debug - Added detection {bbox_idx}: bbox={bbox}, conf={confidence:.2f}")
                    
                    # Also try alternative structures
                    if 'bboxes' in results_dict:
                        print("🔍 Debug - Found direct bboxes in JSON")
                        bboxes = results_dict['bboxes']
                        for bbox_info in bboxes:
                            bbox = bbox_info.get('bbox')
                            confidence = bbox_info.get('confidence', 0.9)
                            if bbox and confidence >= confidence_threshold:
                                area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                                ocr_results.append({
                                    'text': '',
                                    'bbox': bbox,
                                    'confidence': confidence,
                                    'area': area,
                                    'source': 'surya_detection'
                                })
                        
                except Exception as e:
                    print(f"⚠️ Failed to parse detection JSON: {e}")
                    print(f"🔍 Debug - JSON sample: {json_output[:300]}...")
            
            print(f"✅ Surya Detection found {len(ocr_results)} text regions")
            return ocr_results
            
        except Exception as e:
            print(f"❌ Surya Detection failed: {e}")
            import traceback
            traceback.print_exc()
            return []

    def _integrate_basic_surya_ocr_with_text(self, image: Image.Image, confidence_threshold: float = 0.3) -> List[Dict]:
        """Use basic_surya for both detection and recognition to get actual text content"""
        try:
            # Import the basic surya node
            from .basic_surya import SuryaOCRNode
            
            # Initialize if needed
            if not hasattr(self, 'basic_surya_node'):
                self.basic_surya_node = SuryaOCRNode()
            
            # Convert PIL to tensor format expected by basic_surya
            img_array = np.array(image).astype(np.float32) / 255.0
            img_tensor = torch.from_numpy(img_array).unsqueeze(0)
            
            # STRATEGY: First get detections, then run OCR for text content
            print("🔍 Strategy: Detection first, then OCR for text content...")
            
            # Step 1: Get high-quality detections
            det_text_output, det_json_output, det_annotated_image, detection_data = self.basic_surya_node.process_ocr(
                img_tensor, 
                task_mode="detection_only",  # Get the best detections first
                confidence_threshold=confidence_threshold
            )
            
            print(f"🔍 Detection step: found {len(detection_data) if detection_data else 0} text regions")
            
            # Step 2: Run full OCR for text recognition
            ocr_text_output, ocr_json_output, ocr_annotated_image, ocr_data = self.basic_surya_node.process_ocr(
                img_tensor, 
                task_mode="ocr_with_boxes",  # Get text content
                confidence_threshold=confidence_threshold
            )
            
            print(f"🔍 OCR step: text_output={bool(ocr_text_output)}, json_output={bool(ocr_json_output)}")
            
            # Step 3: Parse the OCR results to get text content
            ocr_results_with_text = []
            
            # Try to parse OCR JSON first
            if ocr_json_output:
                import json
                try:
                    ocr_dict = json.loads(ocr_json_output)
                    print(f"🔍 OCR JSON structure: {list(ocr_dict.keys())}")
                    
                    detections = ocr_dict.get('detections', [])
                    if detections and len(detections) > 0:
                        page_detections = detections[0]  # First page
                        if isinstance(page_detections, dict):
                            bboxes = page_detections.get('bboxes', [])
                            print(f"🔍 OCR found {len(bboxes)} bboxes with text")
                            
                            for bbox_info in bboxes:
                                bbox = bbox_info.get('bbox')
                                confidence = bbox_info.get('confidence', 0.9)
                                text = bbox_info.get('text', '').strip()
                                
                                if bbox and confidence >= confidence_threshold and text:
                                    area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                                    ocr_results_with_text.append({
                                        'text': text,
                                        'bbox': bbox,
                                        'confidence': confidence,
                                        'area': area,
                                        'source': 'surya_ocr'
                                    })
                                    
                except Exception as e:
                    print(f"⚠️ Failed to parse OCR JSON: {e}")
            
            # Step 4: If OCR parsing failed, use detection bboxes and extract text manually
            if not ocr_results_with_text and detection_data:
                print("🔧 OCR parsing failed, using detection bboxes + manual text extraction...")
                
                # Parse detection data to get high-quality bounding boxes
                detection_bboxes = []
                if isinstance(detection_data, list):
                    for detection in detection_data:
                        bbox = detection.get('bbox')
                        confidence = detection.get('confidence', 0.9)
                        if bbox and confidence >= confidence_threshold:
                            detection_bboxes.append((bbox, confidence))
                
                elif det_json_output:
                    # Parse detection JSON
                    import json
                    try:
                        det_dict = json.loads(det_json_output)
                        detections = det_dict.get('detections', [])
                        if detections and len(detections) > 0:
                            page_detections = detections[0]
                            if isinstance(page_detections, dict):
                                bboxes = page_detections.get('bboxes', [])
                                for bbox_info in bboxes:
                                    bbox = bbox_info.get('bbox')
                                    confidence = bbox_info.get('confidence', 0.9)
                                    if bbox and confidence >= confidence_threshold:
                                        detection_bboxes.append((bbox, confidence))
                    except Exception as e:
                        print(f"⚠️ Failed to parse detection JSON: {e}")
                
                print(f"🔍 Found {len(detection_bboxes)} detection bboxes for text extraction")
                
                # Use the text from full OCR output if available
                if ocr_text_output:
                    # Split text content and assign to detection boxes
                    text_lines = [line.strip() for line in ocr_text_output.strip().split('\n') if line.strip()]
                    print(f"🔍 Found {len(text_lines)} text lines to match with {len(detection_bboxes)} boxes")
                    
                    # Match text lines to detection boxes (simple approach)
                    for i, (bbox, confidence) in enumerate(detection_bboxes):
                        text = text_lines[i] if i < len(text_lines) else f"Text region {i+1}"
                        area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                        
                        ocr_results_with_text.append({
                            'text': text,
                            'bbox': bbox,
                            'confidence': confidence,
                            'area': area,
                            'source': 'surya_ocr_hybrid'
                        })
                else:
                    # No text content available, use detection boxes with placeholder text
                    print("🔧 No text content available, using detection boxes only...")
                    for i, (bbox, confidence) in enumerate(detection_bboxes):
                        area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
                        ocr_results_with_text.append({
                            'text': f"[Text region {i+1}]",  # Placeholder text
                            'bbox': bbox,
                            'confidence': confidence,
                            'area': area,
                            'source': 'surya_detection_only'
                        })
            
            # Step 5: Final fallback - use just the text output with fake bboxes
            if not ocr_results_with_text and ocr_text_output:
                print("🔧 Final fallback: using text output with estimated bboxes...")
                lines = ocr_text_output.strip().split('\n')
                line_height = image.height // max(len(lines), 1)
                
                for i, line in enumerate(lines):
                    if line.strip():
                        y1 = i * line_height
                        y2 = (i + 1) * line_height
                        ocr_results_with_text.append({
                            'text': line.strip(),
                            'bbox': [50, y1, image.width - 50, y2],
                            'confidence': 0.8,
                            'area': (image.width - 100) * line_height,
                            'source': 'surya_ocr_fallback'
                        })
            
            print(f"✅ Surya OCR found {len(ocr_results_with_text)} text regions with content")
            
            # Debug first few results
            for i, result in enumerate(ocr_results_with_text[:3]):
                print(f"   Text {i}: '{result['text'][:50]}...' conf={result['confidence']:.2f}")
            
            return ocr_results_with_text
            
        except Exception as e:
            print(f"❌ Surya OCR with text failed: {e}")
            import traceback
            traceback.print_exc()
            return []


    def analyze_layout(self, image, primary_ocr="surya", enable_florence2=True, 
                      enable_layoutlmv3=True, enable_cv2_backup=False,
                      extract_images=True, extract_text=True,
                      fusion_strategy="aggressive", florence2_image_prompt=None, florence2_text_prompt=None,
                      min_image_area=10000, min_text_area=500, 
                      surya_confidence=0.3, florence2_confidence=0.5, cv2_confidence=0.7,
                      text_grouping_mode="florence2_priority", text_merge_distance=25,
                      include_text_recognition=True, enable_debug_logging=True,
                      surya_recognition_batch=768, surya_detector_batch=54,
                      smart_grouping_enabled=True, vertical_merge_threshold=15, horizontal_merge_threshold=10):
        """Main analysis function with enhanced text grouping controls"""
        
        # Convert tensor to PIL
        pil_image = tensor_to_PIL(image)
        self.current_image = pil_image
        
        # Store grouping parameters for use in fusion
        self.text_merge_distance = text_merge_distance
        self.smart_grouping_enabled = smart_grouping_enabled
        self.vertical_merge_threshold = vertical_merge_threshold
        self.horizontal_merge_threshold = horizontal_merge_threshold
        
        if enable_debug_logging:
            print(f"🔍 Enhanced Layout Analysis starting:")
            print(f"  Image size: {pil_image.size}")
            print(f"  Primary OCR: {primary_ocr}")
            print(f"  Florence2: {enable_florence2}, LayoutLMv3: {enable_layoutlmv3}")
            print(f"  Text grouping: {text_grouping_mode}, merge_distance: {text_merge_distance}")
            print(f"  Text recognition: {include_text_recognition}")
            print(f"  Smart grouping: {smart_grouping_enabled}, v_thresh: {vertical_merge_threshold}, h_thresh: {horizontal_merge_threshold}")
        
        # Initialize components as needed
        ocr_results = []
        florence2_images = []
        florence2_text = []
        cv2_images = []
        cv2_text = []
        
        # 1. RUN SURYA OCR (with improved strategy)
        if primary_ocr == "surya":
            if include_text_recognition:
                # Full OCR with enhanced text extraction strategy
                ocr_results = self._integrate_basic_surya_ocr_with_text(pil_image, surya_confidence)
            else:
                # Detection only for better coverage (existing working code)
                ocr_results = self._integrate_basic_surya_detection(pil_image, surya_confidence)
            
            # Fallback to other OCR if Surya fails
            if not ocr_results:
                print("🔧 Surya failed, trying PaddleOCR fallback...")
                if not hasattr(self, 'fallback_ocr') or self.fallback_ocr is None:
                    self._init_fallback_ocr("paddleocr")
                
                if hasattr(self, 'fallback_ocr') and self.fallback_ocr:
                    ocr_results = self._run_fallback_ocr(pil_image, "paddleocr")
        
        # Fallback to PaddleOCR if requested
        elif primary_ocr in ["paddleocr", "easyocr"]:
            if not hasattr(self, 'fallback_ocr') or self.fallback_ocr is None:
                self._init_fallback_ocr(primary_ocr)
            
            if hasattr(self, 'fallback_ocr') and self.fallback_ocr:
                ocr_results = self._run_fallback_ocr(pil_image, primary_ocr)
        
        # 2. RUN FLORENCE2 DETECTION with optimized prompts
        if enable_florence2:
            if self.florence2_detector is None:
                self._init_florence2()
            
            if self.florence2_detector:
                florence2_images, florence2_text = self._run_florence2_detection(
                    pil_image,
                    image_prompt=florence2_image_prompt or "rectangular images in page",
                    text_prompt=florence2_text_prompt or "text, caption, paragraph, title"
                )
        
        # 3. RUN CV2 ANALYSIS only if specifically requested (minimal use)
        if enable_cv2_backup:
            cv2_images, cv2_text = self._run_cv2_analysis(pil_image)
        
        # 4. ENHANCED FUSION with LayoutLMv3
        final_images, final_text = self._smart_fusion_with_layoutlmv3(
            ocr_results, florence2_images, florence2_text, cv2_images, cv2_text,
            fusion_strategy, min_image_area, min_text_area,
            surya_confidence, florence2_confidence, cv2_confidence, text_grouping_mode,
            enable_layoutlmv3
        )

        # 5. CREATE OUTPUTS
        image_boxes = final_images if extract_images else []
        text_boxes = final_text if extract_text else []
        
        # Create visualization
        overlay = self._create_enhanced_overlay(pil_image, image_boxes, text_boxes)
        overlay_tensor = PIL_to_tensor(overlay)
        
        # Extract crops
        image_crops = []
        text_crops = []
        
        if extract_images:
            for img_box in image_boxes:
                try:
                    crop = pil_image.crop(img_box['bbox'])
                    image_crops.append(PIL_to_tensor(crop))
                except Exception as e:
                    print(f"⚠️ Error cropping image box: {e}")
        
        if extract_text:
            for txt_box in text_boxes:
                try:
                    crop = pil_image.crop(txt_box['bbox'])
                    text_crops.append(PIL_to_tensor(crop))
                except Exception as e:
                    print(f"⚠️ Error cropping text box: {e}")
        
        # 6. EXTRACT ALL TEXT CONTENT
        extracted_text = ""
        if include_text_recognition and text_boxes:
            text_content = []
            for txt_box in text_boxes:
                if txt_box.get('text', '').strip():
                    text_content.append(txt_box['text'].strip())
            
            extracted_text = "\n".join(text_content)
            print(f"📝 Extracted {len(text_content)} text blocks, total length: {len(extracted_text)} characters")
        
        # Analysis results
        analysis_results = {
            "total_detections": len(image_boxes) + len(text_boxes),
            "image_detections": len(image_boxes),
            "text_detections": len(text_boxes),
            "fusion_strategy": fusion_strategy,
            "text_grouping_mode": text_grouping_mode,
            "text_recognition_enabled": include_text_recognition,
            "extracted_text_length": len(extracted_text),
            "layoutlmv3_enabled": enable_layoutlmv3,
            "confidence_thresholds": {
                "surya": surya_confidence,
                "florence2": florence2_confidence,
                "cv2": cv2_confidence
            },
            "detection_methods_used": {
                "ocr": len(ocr_results) > 0,
                "florence2": len(florence2_images) + len(florence2_text) > 0,
                "cv2": len(cv2_images) + len(cv2_text) > 0,
                "layoutlmv3": enable_layoutlmv3
            },
            "method_statistics": {
                "ocr_results": len(ocr_results),
                "florence2_images": len(florence2_images),
                "florence2_text": len(florence2_text),
                "cv2_images": len(cv2_images),
                "cv2_text": len(cv2_text)
            }
        }
        
        return (overlay_tensor, image_boxes, text_boxes, image_crops, text_crops, extracted_text, analysis_results)

    
    def _create_enhanced_overlay(self, image: Image.Image, image_boxes: List[Dict], text_boxes: List[Dict]) -> Image.Image:
        """Create enhanced visualization showing detection sources"""
        overlay = image.copy()
        draw = ImageDraw.Draw(overlay)
        
        # Color coding by detection method
        colors = {
            'florence2': (255, 0, 0),        # Red
            'surya': (0, 255, 0),            # Green  
            'surya_detection': (0, 200, 0),  # Dark Green
            'surya_ocr': (0, 255, 50),       # Bright Green
            'surya_grouped': (0, 180, 100),  # Grouped Green
            'paddleocr': (0, 150, 255),      # Light Blue
            'easyocr': (100, 150, 255),      # Lighter Blue
            'cv2': (0, 0, 255),              # Blue
            'layoutlmv3': (255, 0, 255),     # Magenta
            'layoutlmv3_semantic': (200, 0, 255),  # Purple
            'multi': (255, 165, 0)           # Orange for multi-method
        }
        
        # Draw image boxes
        for i, box in enumerate(image_boxes):
            # Convert bbox to integers
            bbox = [int(coord) for coord in box['bbox']]
            methods = box['detection_methods']
            
            # Choose color based on detection methods
            if len(methods) > 1:
                color = colors['multi']
                method_label = "+".join(methods)
            else:
                color = colors.get(methods[0], (128, 128, 128))
                method_label = methods[0]
            
            # Draw rectangle
            draw.rectangle(bbox, outline=color, width=3)
            
            # Add label with semantic info if available
            label = f"IMG{i+1}: {method_label} ({box['confidence']:.2f})"
            if box.get('semantic_label'):
                label += f" [{box['semantic_label']}]"
            draw.text((bbox[0], max(0, bbox[1] - 20)), label, fill=color)
        
        # Draw text boxes
        for i, box in enumerate(text_boxes):
            # Convert bbox to integers
            bbox = [int(coord) for coord in box['bbox']]
            methods = box['detection_methods']
            
            # Choose color based on detection methods
            if len(methods) > 1:
                color = colors['multi']
                method_label = "+".join(methods)
            else:
                color = colors.get(methods[0], (128, 128, 128))
                method_label = methods[0]
            
            # Draw rectangle (dashed style for text)
            self._draw_dashed_rectangle(draw, bbox, color, width=2)
            
            # Add label with semantic info if available
            label = f"TXT{i+1}: {method_label} ({box['confidence']:.2f})"
            if box.get('semantic_label'):
                label += f" [{box['semantic_label']}]"
            draw.text((bbox[0], max(0, bbox[1] - 20)), label, fill=color)
        
        return overlay
    
    def _draw_dashed_rectangle(self, draw, bbox, color, width=2, dash_length=5):
        """Draw a dashed rectangle"""
        # Convert all coordinates to integers
        x1, y1, x2, y2 = map(int, bbox)
        
        # Top edge
        for x in range(x1, x2, dash_length * 2):
            draw.line([(x, y1), (min(x + dash_length, x2), y1)], fill=color, width=width)
        
        # Bottom edge  
        for x in range(x1, x2, dash_length * 2):
            draw.line([(x, y2), (min(x + dash_length, x2), y2)], fill=color, width=width)
        
        # Left edge
        for y in range(y1, y2, dash_length * 2):
            draw.line([(x1, y), (x1, min(y + dash_length, y2))], fill=color, width=width)
        
        # Right edge
        for y in range(y1, y2, dash_length * 2):
            draw.line([(x2, y), (x2, min(y + dash_length, y2))], fill=color, width=width)



# Node mappings
NODE_CLASS_MAPPINGS = {
    "EnhancedLayoutParser_v02": EnhancedLayoutParserNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "EnhancedLayoutParser_v02": "Enhanced Layout Parser (Multi-Modal) v02",
}