Update working_yolo_pipeline.py

working_yolo_pipeline.py

@@ -2131,10 +2131,26 @@ def run_document_pipeline(input_path: str, layoutlmv3_model_path: str):
             print("❌ No data extracted.")
             return None
 
+        # # 3. CONSOLIDATE BLOCKS FOR INFERENCE
+        # sequential_blocks = []
+        # for p_data in all_pages_data:
+        #     sequential_blocks.extend(p_data.get('blocks', []))
+
         # 3. CONSOLIDATE BLOCKS FOR INFERENCE
         sequential_blocks = []
         for p_data in all_pages_data:
-            sequential_blocks.extend(p_data.get('blocks', []))
+            if isinstance(p_data, dict):
+                # If it's a dictionary, extract the 'blocks' key
+                blocks = p_data.get('blocks', [])
+                sequential_blocks.extend(blocks)
+            elif isinstance(p_data, list):
+                # If it's already a list, add it directly
+                sequential_blocks.extend(p_data)
+            else:
+                print(f"⚠️ Warning: Unexpected data type in all_pages_data: {type(p_data)}")
+
+
+                
 
         # --- 4. STARTING LAYOUTLMV3 INFERENCE ---
         print("\n" + "=" * 80)