Update pdf_section_extractor.py

pdf_section_extractor.py

@@ -1,6 +1,223 @@
+# import pdfplumber
+# import pandas as pd
+# import re
+# from typing import List, Dict, Tuple, Any
+
+# class PDFSectionExtractor:
+#     def __init__(self, pdf_path: str):
+#         """Initialize with path to PDF file."""
+#         self.pdf_path = pdf_path
+#         self.tables = []
+#         self.table_names = {}  # Store table names and their content
+        
+#     def is_level_three_section(self, section_number: str) -> bool:
+#         """Check if the section number is a level three section (e.g., 2.1.1)."""
+#         return len(section_number.split('.')) == 3
+
+#     def find_table_names(self, text: str) -> List[Dict[str, str]]:
+#         """Extract table names from text."""
+#         table_pattern = r'\*\*Table\s+(\d+):\s+([^*]+)\*\*'
+#         return [(match.group(1), match.group(2).strip()) 
+#                 for match in re.finditer(table_pattern, text)]
+
+#     def clean_content(self, content: str) -> str:
+#         """Remove table references and names from content."""
+#         # Remove table references
+#         content = re.sub(r'\*\*Table\s+\d+:\s+[^*]+\*\*', '', content)
+#         # Remove any empty lines created
+#         content = '\n'.join(line for line in content.split('\n') if line.strip())
+#         return content
+        
+#     def merge_split_tables(self, tables: List[List]) -> List[List]:
+#         """Merge tables that are split across pages."""
+#         merged_tables = []
+#         current_table = None
+        
+#         for table in tables:
+#             if not table:
+#                 continue
+                
+#             if current_table is None:
+#                 current_table = table
+#             else:
+#                 # Check if this table is a continuation
+#                 # Compare the number of columns
+#                 if len(table[0]) == len(current_table[0]):
+#                     current_table.extend(table)
+#                 else:
+#                     merged_tables.append(current_table)
+#                     current_table = table
+        
+#         if current_table:
+#             merged_tables.append(current_table)
+            
+#         return merged_tables
+
+#     def extract_tables(self) -> List[Dict[str, Any]]:
+#         """Extract tables with their names and merge split tables."""
+#         tables_data = []
+#         current_section = None
+#         current_table_data = None
+#         current_table_name = None
+        
+#         with pdfplumber.open(self.pdf_path) as pdf:
+#             for page in pdf.pages:
+#                 text = page.extract_text(x_tolerance=1,y_tolerance=0) or ''
+                
+#                 # Find table names in the text
+#                 table_names = self.find_table_names(text)
+#                 tables = page.extract_tables()
+                
+#                 # Process each table found
+#                 if tables:
+#                     tables = self.merge_split_tables(tables)
+                    
+#                     for i, table in enumerate(tables):
+#                         table_name = None
+#                         if i < len(table_names):
+#                             table_num, name = table_names[i]
+#                             table_name = f"Table {table_num}: {name}"
+                        
+#                         if table:  # Check if table has content
+#                             df = pd.DataFrame(table)
+#                             # Clean the DataFrame
+#                             df = df.dropna(how='all').dropna(axis=1, how='all')
+#                             # Replace None with empty string
+#                             df = df.fillna('')
+                            
+#                             tables_data.append({
+#                                 'name': table_name,
+#                                 'data': df
+#                             })
+        
+#         return tables_data
+
+#     def extract_sections(self) -> List[Dict[str, str]]:
+#         """Extract sections from PDF with content, excluding tables."""
+#         sections = []
+#         current_section = None
+#         current_content = []
+#         section_pattern = r'^(\d+\.(?:\d+)?(?:\.\d+)?)\s+(.+)$'
+        
+#         with pdfplumber.open(self.pdf_path) as pdf:
+#             for page in pdf.pages:
+#                 text = page.extract_text(x_tolerance=1)
+#                 if not text:
+#                     continue
+                    
+#                 lines = text.split('\n')
+                
+#                 for line in lines:
+#                     match = re.match(section_pattern, line.strip())
+                    
+#                     if match:
+#                         if current_section:
+#                             content_text = '\n'.join(current_content)
+#                             content_text = self.clean_content(content_text)
+                            
+#                             if self.is_level_three_section(current_section[0]):
+#                                 full_content = current_section[1] + '\n' + content_text
+#                                 sections.append({
+#                                     'section_number': current_section[0],
+#                                     'section_name': '',
+#                                     'content': full_content
+#                                 })
+#                             else:
+#                                 sections.append({
+#                                     'section_number': current_section[0],
+#                                     'section_name': current_section[1],
+#                                     'content': content_text
+#                                 })
+                        
+#                         current_section = (match.group(1), match.group(2))
+#                         current_content = []
+#                     elif current_section:
+#                         current_content.append(line.strip())
+        
+#         # Handle the last section
+#         if current_section:
+#             content_text = '\n'.join(current_content)
+#             content_text = self.clean_content(content_text)
+            
+#             if self.is_level_three_section(current_section[0]):
+#                 full_content = current_section[1] + '\n' + content_text
+#                 sections.append({
+#                     'section_number': current_section[0],
+#                     'section_name': '',
+#                     'content': full_content
+#                 })
+#             else:
+#                 sections.append({
+#                     'section_number': current_section[0],
+#                     'section_name': current_section[1],
+#                     'content': content_text
+#                 })
+                
+#         return sections
+
+# def convert_pdf_to_excel(pdf_path: str, excel_path: str):
+#     """Convert PDF with sections and tables to Excel file."""
+#     try:
+#         extractor = PDFSectionExtractor(pdf_path)
+#         sections = extractor.extract_sections()
+#         tables = extractor.extract_tables()
+        
+#         with pd.ExcelWriter(excel_path, engine='openpyxl') as writer:
+#             # Write sections to main sheet
+#             df_sections = pd.DataFrame(sections)
+#             df_sections.to_excel(writer, index=False, sheet_name='Sections')
+            
+#             # Auto-adjust sections sheet
+#             worksheet = writer.sheets['Sections']
+#             for idx, col in enumerate(['A', 'B', 'C']):
+#                 worksheet.column_dimensions[col].width = 15 if idx < 2 else 50
+            
+#             # Write tables to separate sheets
+#             for i, table_info in enumerate(tables, 1):
+#                 if table_info['name']:
+#                     sheet_name = table_info['name'][:31]  # Excel sheet name length limit
+#                 else:
+#                     sheet_name = f'Table_{i}'
+                
+#                 # Write table data
+#                 table_info['data'].to_excel(writer, sheet_name=sheet_name, index=False)
+                
+#                 # Auto-adjust table sheet
+#                 worksheet = writer.sheets[sheet_name]
+#                 for column in worksheet.columns:
+#                     max_length = 0
+#                     column = [cell for cell in column]
+#                     for cell in column:
+#                         try:
+#                             if len(str(cell.value)) > max_length:
+#                                 max_length = len(cell.value)
+#                         except:
+#                             pass
+#                     adjusted_width = (max_length + 2)
+#                     worksheet.column_dimensions[column[0].column_letter].width = adjusted_width
+        
+#         return True
+    
+#     except Exception as e:
+#         print(f"Error converting PDF to Excel: {str(e)}")
+#         return False
+
+# if __name__ == "__main__":
+#     pdf_path = "/Users/aakanksha.n/Desktop/pdf_to_excel/[00 12 10] 26251-100-3DR-S04-00001_002.docx.pdf"
+#     excel_path = "/Users/aakanksha.n/Desktop/pdf_to_excel/format4.xlsx"
+    
+#     success = convert_pdf_to_excel(pdf_path, excel_path)
+#     if success:
+#         print("Successfully converted PDF to Excel!")
+#     else:
+#         print("Failed to convert PDF to Excel.")
+
+
+
 import pdfplumber
 import pandas as pd
 import re
+import os
 from typing import List, Dict, Tuple, Any
 
 class PDFSectionExtractor:
@@ -8,8 +225,63 @@ class PDFSectionExtractor:
         """Initialize with path to PDF file."""
         self.pdf_path = pdf_path
         self.tables = []
-        self.table_names = {}  # Store table names and their content
-        
+        self.table_names = {}
+        self.table_content_markers = set()
+        self.table_content_lines = set()
+        self.table_headers = set()
+        self.document_name = None  # Set to None to identify when it's missing
+        self.document_id = None
+
+    def extract_document_info_from_pdf(self) -> Tuple[str, str]:
+        """Extract document name and ID from the first page of PDF."""
+        with pdfplumber.open(self.pdf_path) as pdf:
+            if not pdf.pages:
+                return "", ""
+
+            first_page_text = pdf.pages[0].extract_text()
+
+            if not first_page_text:
+                return "", ""
+
+            doc_name_pattern = r'[dD]ocument\s*[nN]ame:\s*(.*?)(?:\n|$)'
+            doc_id_pattern = r'[dD]ocument\s*[iI][dD]:\s*(.*?)(?:\n|$)'
+
+
+            doc_name_match = re.search(doc_name_pattern, first_page_text)
+            doc_id_match = re.search(doc_id_pattern, first_page_text)
+            print(f" extract document info function{doc_name_match}")
+
+            doc_name = doc_name_match.group(1).strip() if doc_name_match else ""
+            doc_id = doc_id_match.group(1).strip() if doc_id_match else ""
+
+            # print(f"Document name from pdf inside function{doc_name},{doc_id}")
+
+            return doc_name, doc_id
+
+    def get_document_info(self) -> Tuple[str, str]:
+        """
+        Get document info from the PDF content.
+        If missing, fall back to filename for document ID only.
+        """
+        if self.document_name is None or self.document_id is None:
+            # Try extracting from PDF
+            doc_name, doc_id = self.extract_document_info_from_pdf()
+            self.document_name = doc_name
+            self.document_id = doc_id
+            # print(f"Entering if one,{self.document_name}")
+
+        # Fallback only if document ID or name is missing
+        if not self.document_name or not self.document_id:
+            doc_name = os.path.basename(self.pdf_path)
+            doc_id_match = re.search(r'\[\d{2}\s+\d{2}\s+\d{2}\]\s+(\d+)', doc_name)
+            if not self.document_id:
+                self.document_id = doc_id_match.group(1) if doc_id_match else ""
+            if not self.document_name:
+                self.document_name = doc_name
+
+        # print(f"final{self.document_name},{self.document_id}")
+        return self.document_name, self.document_id
+
     def is_level_three_section(self, section_number: str) -> bool:
         """Check if the section number is a level three section (e.g., 2.1.1)."""
         return len(section_number.split('.')) == 3
@@ -20,77 +292,139 @@ class PDFSectionExtractor:
         return [(match.group(1), match.group(2).strip()) 
                 for match in re.finditer(table_pattern, text)]
 
-    def clean_content(self, content: str) -> str:
-        """Remove table references and names from content."""
-        # Remove table references
-        content = re.sub(r'\*\*Table\s+\d+:\s+[^*]+\*\*', '', content)
-        # Remove any empty lines created
-        content = '\n'.join(line for line in content.split('\n') if line.strip())
-        return content
-        
-    def merge_split_tables(self, tables: List[List]) -> List[List]:
-        """Merge tables that are split across pages."""
-        merged_tables = []
-        current_table = None
-        
-        for table in tables:
-            if not table:
-                continue
-                
-            if current_table is None:
-                current_table = table
+    def store_table_headers(self, table: List[List[str]]):
+        """Store table headers."""
+        if table and table[0]:
+            header_row = table[0]
+            for header in header_row:
+                if header:
+                    header_text = self.fix_table_cell_spacing(header)
+                    if header_text:
+                        self.table_headers.add(header_text)
+
+    def store_table_content(self, table: List[List[str]]):
+        """Store all content from table cells, including headers."""
+        if not table:
+            return
+
+        # Store headers separately
+        self.store_table_headers(table)
+
+        # Store all cell content
+        for row in table:
+            for cell in row:
+                if cell:
+                    cell_text = self.fix_table_cell_spacing(cell)
+                    if cell_text:
+                        self.table_content_lines.add(cell_text)
+
+    def fix_table_cell_spacing(self, cell_text: str) -> str:
+        """Fix spacing issues within table cells."""
+        if not isinstance(cell_text, str):
+            return str(cell_text)
+
+        text = str(cell_text).strip()
+        text = re.sub(r',(?=\S)', ', ', text)
+        text = re.sub(r'(\d)([a-zA-Z])', r'\1 \2', text)
+        text = re.sub(r'([a-zA-Z])(\d)', r'\1 \2', text)
+        text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text)
+        text = re.sub(r'\s+', ' ', text)
+        return text.strip()
+
+    def contains_table_content(self, line: str) -> bool:
+        """Check if a line contains any stored table content or headers."""
+        line = self.fix_table_cell_spacing(line)
+        return (any(table_line in line for table_line in self.table_content_lines) or
+                any(header in line for header in self.table_headers))
+
+    def is_table_content(self, line: str) -> bool:
+        """Identify if a line is part of a table."""
+        if any(marker in line for marker in ['|', '+', '─', '│', '��', '┐', '└', '┘', '├', '┤', '┬', '┴']):
+            return True
+        if re.search(r'\S+\s{2,}\S+', line):
+            return True
+        if re.match(r'^[\s\-+|=_]{3,}$', line):
+            return True
+        return False
+
+    def filter_table_content(self, content: str) -> str:
+        """Remove lines that contain table content."""
+        lines = content.split('\n')
+        filtered_lines = []
+
+        for line in lines:
+            line = line.strip()
+            if line and not self.contains_table_content(line):
+                filtered_lines.append(line)
+
+        return '\n'.join(filtered_lines)
+
+    def process_table(self, table: List[List[str]], table_num: int) -> pd.DataFrame:
+        """Process table data with proper spacing and ensure unique column names."""
+        if not table:
+            return pd.DataFrame()
+
+        # Store all table content including headers
+        self.store_table_content(table)
+
+        # Create DataFrame with default column names
+        df = pd.DataFrame(table)
+        if df.empty:
+            return df
+
+        # Get the header row (first row)
+        header = df.iloc[0].tolist()
+
+        # Create unique column names
+        unique_cols = []
+        col_count = {}
+
+        for col in header:
+            col = str(col) if col else "Unnamed"
+            if col in col_count:
+                col_count[col] += 1
+                unique_cols.append(f"{col}_{col_count[col]}")
             else:
-                # Check if this table is a continuation
-                # Compare the number of columns
-                if len(table[0]) == len(current_table[0]):
-                    current_table.extend(table)
-                else:
-                    merged_tables.append(current_table)
-                    current_table = table
-        
-        if current_table:
-            merged_tables.append(current_table)
-            
-        return merged_tables
+                col_count[col] = 0
+                unique_cols.append(col)
+
+        # Set the unique column names
+        df.columns = unique_cols
+
+        # Remove the header row since we used it for column names
+        df = df.iloc[1:].reset_index(drop=True)
+
+        # Fix spacing in cells
+        for col in df.columns:
+            df[col] = df[col].astype(str).apply(self.fix_table_cell_spacing)
+
+        # Add metadata
+        df.insert(0, 'Document_Name', self.document_name)
+        df.insert(1, 'Document_ID', self.document_id)
+        df.insert(2, 'Table_Number', table_num)
+
+        return df
+
+    def extract_tables(self) -> List[pd.DataFrame]:
+        """Extract tables from PDF and return a list of DataFrames."""
+        tables = []
+        table_counter = 1
+
+        # Ensure we have document info before processing tables
+        self.get_document_info()
 
-    def extract_tables(self) -> List[Dict[str, Any]]:
-        """Extract tables with their names and merge split tables."""
-        tables_data = []
-        current_section = None
-        current_table_data = None
-        current_table_name = None
-        
         with pdfplumber.open(self.pdf_path) as pdf:
             for page in pdf.pages:
-                text = page.extract_text(x_tolerance=1,y_tolerance=0) or ''
-                
-                # Find table names in the text
-                table_names = self.find_table_names(text)
-                tables = page.extract_tables()
-                
-                # Process each table found
-                if tables:
-                    tables = self.merge_split_tables(tables)
-                    
-                    for i, table in enumerate(tables):
-                        table_name = None
-                        if i < len(table_names):
-                            table_num, name = table_names[i]
-                            table_name = f"Table {table_num}: {name}"
-                        
-                        if table:  # Check if table has content
-                            df = pd.DataFrame(table)
-                            # Clean the DataFrame
-                            df = df.dropna(how='all').dropna(axis=1, how='all')
-                            # Replace None with empty string
-                            df = df.fillna('')
-                            
-                            tables_data.append({
-                                'name': table_name,
-                                'data': df
-                            })
-        
-        return tables_data
+                detected_tables = page.extract_tables()
+
+                for table in detected_tables:
+                    if table:
+                        df = self.process_table(table, table_counter)
+                        if not df.empty:
+                            tables.append(df)
+                            table_counter += 1
+
+        return tables
 
     def extract_sections(self) -> List[Dict[str, str]]:
         """Extract sections from PDF with content, excluding tables."""
@@ -98,116 +432,120 @@ class PDFSectionExtractor:
         current_section = None
         current_content = []
         section_pattern = r'^(\d+\.(?:\d+)?(?:\.\d+)?)\s+(.+)$'
-        
+
+        # Ensure we have document info before processing sections
+        doc_name, doc_id = self.get_document_info()
+        table_counter = 1
+
         with pdfplumber.open(self.pdf_path) as pdf:
             for page in pdf.pages:
                 text = page.extract_text(x_tolerance=1)
                 if not text:
                     continue
-                    
+
                 lines = text.split('\n')
-                
+                table_found = False
+
                 for line in lines:
-                    match = re.match(section_pattern, line.strip())
-                    
+                    line = line.strip()
+                    match = re.match(section_pattern, line)
+
                     if match:
                         if current_section:
-                            content_text = '\n'.join(current_content)
-                            content_text = self.clean_content(content_text)
-                            
-                            if self.is_level_three_section(current_section[0]):
-                                full_content = current_section[1] + '\n' + content_text
-                                sections.append({
-                                    'section_number': current_section[0],
-                                    'section_name': '',
-                                    'content': full_content
-                                })
-                            else:
-                                sections.append({
-                                    'section_number': current_section[0],
-                                    'section_name': current_section[1],
-                                    'content': content_text
-                                })
-                        
+                            content_text = self.process_content(current_content)
+                            content_text = self.filter_table_content(content_text)
+                            sections.append(self.create_section_dict(
+                                doc_name, doc_id, current_section, content_text))
+
                         current_section = (match.group(1), match.group(2))
                         current_content = []
                     elif current_section:
-                        current_content.append(line.strip())
-        
-        # Handle the last section
+                        if self.is_table_content(line):
+                            if not table_found:
+                                current_content.append(f"[Table {table_counter} is available in Table_{table_counter} sheet]")
+                                table_counter += 1
+                                table_found = True
+                        else:
+                            current_content.append(line)
+                            table_found = False
+
         if current_section:
-            content_text = '\n'.join(current_content)
-            content_text = self.clean_content(content_text)
-            
-            if self.is_level_three_section(current_section[0]):
-                full_content = current_section[1] + '\n' + content_text
-                sections.append({
-                    'section_number': current_section[0],
-                    'section_name': '',
-                    'content': full_content
-                })
-            else:
-                sections.append({
-                    'section_number': current_section[0],
-                    'section_name': current_section[1],
-                    'content': content_text
-                })
-                
+            content_text = self.process_content(current_content)
+            content_text = self.filter_table_content(content_text)
+            sections.append(self.create_section_dict(
+                doc_name, doc_id, current_section, content_text))
+
         return sections
 
+    def process_content(self, content_lines: List[str]) -> str:
+        """Process content lines and remove duplicate table references."""
+        processed_lines = []
+        last_line_was_table_ref = False
+
+        for line in content_lines:  # Fixed: using content_lines parameter
+            line = line.strip()
+            if '[Table' in line:
+                if not last_line_was_table_ref:
+                    processed_lines.append(line)
+                    last_line_was_table_ref = True
+            else:
+                processed_lines.append(line)
+                last_line_was_table_ref = False
+
+        return '\n'.join(processed_lines)
+
+    def create_section_dict(self, doc_name: str, doc_id: str, 
+                            section_tuple: Tuple[str, str], content: str) -> Dict[str, str]:
+        """Create a dictionary for section data."""
+        section_number, section_name = section_tuple
+
+        if self.is_level_three_section(section_number):
+            return {
+                'document_name': doc_name,
+                'document_id': doc_id,
+                'section_number': section_number,
+                'section_name': '',
+                'content': f"{section_name}\n{content}"
+            }
+        else:
+            return {
+                'document_name': doc_name,
+                'document_id': doc_id,
+                'section_number': section_number,
+                'section_name': section_name,
+                'content': content
+            }
+
 def convert_pdf_to_excel(pdf_path: str, excel_path: str):
-    """Convert PDF with sections and tables to Excel file."""
+    """Convert PDF with sections and tables to Excel file with separate sheets for each table."""
     try:
         extractor = PDFSectionExtractor(pdf_path)
-        sections = extractor.extract_sections()
+
+        # Extract and process tables
         tables = extractor.extract_tables()
-        
+
+        # Extract sections
+        sections = extractor.extract_sections()
+        df_sections = pd.DataFrame(sections)
+
+        # Write to Excel file
         with pd.ExcelWriter(excel_path, engine='openpyxl') as writer:
-            # Write sections to main sheet
-            df_sections = pd.DataFrame(sections)
+            # Write sections sheet
             df_sections.to_excel(writer, index=False, sheet_name='Sections')
             
-            # Auto-adjust sections sheet
-            worksheet = writer.sheets['Sections']
-            for idx, col in enumerate(['A', 'B', 'C']):
-                worksheet.column_dimensions[col].width = 15 if idx < 2 else 50
-            
-            # Write tables to separate sheets
-            for i, table_info in enumerate(tables, 1):
-                if table_info['name']:
-                    sheet_name = table_info['name'][:31]  # Excel sheet name length limit
-                else:
-                    sheet_name = f'Table_{i}'
-                
-                # Write table data
-                table_info['data'].to_excel(writer, sheet_name=sheet_name, index=False)
-                
-                # Auto-adjust table sheet
-                worksheet = writer.sheets[sheet_name]
-                for column in worksheet.columns:
-                    max_length = 0
-                    column = [cell for cell in column]
-                    for cell in column:
-                        try:
-                            if len(str(cell.value)) > max_length:
-                                max_length = len(cell.value)
-                        except:
-                            pass
-                    adjusted_width = (max_length + 2)
-                    worksheet.column_dimensions[column[0].column_letter].width = adjusted_width
+            # Write each table to its own sheet
+            for i, table_df in enumerate(tables, 1):
+                sheet_name = f'Table_{i}'
+                table_df.to_excel(writer, index=False, sheet_name=sheet_name)
         
-        return True
-    
+        print(f"Successfully saved PDF data to {excel_path}")
+        print(f"Created {len(tables)} table sheets")
     except Exception as e:
-        print(f"Error converting PDF to Excel: {str(e)}")
-        return False
+        print(f"Error occurred: {e}")
+        raise
 
 if __name__ == "__main__":
-    pdf_path = "/Users/aakanksha.n/Desktop/pdf_to_excel/[00 12 10] 26251-100-3DR-S04-00001_002.docx.pdf"
-    excel_path = "/Users/aakanksha.n/Desktop/pdf_to_excel/format4.xlsx"
+    pdf_path = "[00 12 10] 26251-100-3DR-S04-00001_002.docx.pdf"
+    excel_path = "format.xlsx"
     
-    success = convert_pdf_to_excel(pdf_path, excel_path)
-    if success:
-        print("Successfully converted PDF to Excel!")
-    else:
-        print("Failed to convert PDF to Excel.")
+    convert_pdf_to_excel(pdf_path, excel_path)