SafeVixAI/SafeVixAI-Dataset-Hub

1"""2extract_morth2022_tables.py3===========================4Extracts tabular accident data from the raw MoRTH 2022 PDF reports that were5downloaded into chatbot_service/data/accidents/morth_2022/.6 7The morth_2022 folder currently has two large PDFs but only one tabular CSV.8This script uses pdfplumber to extract all tables from those PDFs and saves9them as clean, labelled CSVs — matching the format of morth_2021/ and morth_2020/.10 11Run:12    pip install pdfplumber13    python scripts/extract_morth2022_tables.py14 15Output:16    chatbot_service/data/accidents/morth_2022/extracted_table_*.csv17"""18from __future__ import annotations19 20import csv21import re22import sys23from pathlib import Path24 25try:26    import pdfplumber27except ImportError:28    print("ERROR: pdfplumber not installed. Run: pip install pdfplumber")29    sys.exit(1)30 31PROJECT_ROOT = Path(__file__).resolve().parents[1]32MORTH_2022_DIR = PROJECT_ROOT / "chatbot_service" / "data" / "accidents" / "morth_2022"33 34 35def _clean_cell(text: str | None) -> str:36    """Normalise whitespace in a table cell value."""37    if text is None:38        return ""39    return re.sub(r"\s+", " ", text.strip())40 41 42def _is_empty_row(row: list[str]) -> bool:43    return all(c == "" for c in row)44 45 46def _is_header_row(row: list[str]) -> bool:47    """Heuristic: a row is a header if most cells look like labels not numbers."""48    non_empty = [c for c in row if c]49    if not non_empty:50        return False51    numeric_count = sum(1 for c in non_empty if re.match(r"^[\d,.\s]+$", c))52    return numeric_count < len(non_empty) / 253 54 55def extract_tables_from_pdf(pdf_path: Path, output_dir: Path) -> int:56    """Extract all tables from a PDF and write them to numbered CSVs."""57    output_dir.mkdir(parents=True, exist_ok=True)58    stem = pdf_path.stem[:24]  # keep filename manageable59    tables_written = 060 61    print(f"\nProcessing: {pdf_path.name} ({pdf_path.stat().st_size / 1_048_576:.1f} MB)")62 63    with pdfplumber.open(pdf_path) as pdf:64        global_table_idx = 065        buffer_rows: list[list[str]] = []66        buffer_header: list[str] = []67 68        for page_num, page in enumerate(pdf.pages, start=1):69            tables = page.extract_tables()70            if not tables:71                continue72 73            for table in tables:74                if not table:75                    continue76 77                cleaned = [78                    [_clean_cell(cell) for cell in row]79                    for row in table80                ]81                cleaned = [r for r in cleaned if not _is_empty_row(r)]82 83                if not cleaned:84                    continue85 86                # Detect if this page continues a previous table (no header in first row)87                first_row_looks_like_header = _is_header_row(cleaned[0])88 89                if first_row_looks_like_header and buffer_rows:90                    # Flush previous buffer91                    _write_table(output_dir, stem, global_table_idx, buffer_header, buffer_rows)92                    tables_written += 193                    global_table_idx += 194                    buffer_rows = []95                    buffer_header = []96 97                if first_row_looks_like_header:98                    buffer_header = cleaned[0]99                    buffer_rows = cleaned[1:]100                else:101                    # Continuation of previous table102                    buffer_rows.extend(cleaned)103 104        # Flush any remaining buffer105        if buffer_rows:106            _write_table(output_dir, stem, global_table_idx, buffer_header, buffer_rows)107            tables_written += 1108 109    return tables_written110 111 112def _write_table(113    output_dir: Path,114    stem: str,115    index: int,116    header: list[str],117    rows: list[list[str]],118) -> None:119    filename = output_dir / f"extracted_{stem}_table_{index:03d}.csv"120    with filename.open("w", newline="", encoding="utf-8") as fp:121        writer = csv.writer(fp)122        if header:123            writer.writerow(header)124        writer.writerows(rows)125    print(f"  Wrote: {filename.name} ({len(rows)} data rows)")126 127 128def main() -> None:129    pdfs = sorted(MORTH_2022_DIR.glob("*.pdf"))130    if not pdfs:131        print(f"No PDFs found in {MORTH_2022_DIR}")132        print("Download the MoRTH 2022 report from:")133        print("  https://morth.nic.in/road-accident-in-india")134        sys.exit(1)135 136    total_tables = 0137    for pdf_path in pdfs:138        n = extract_tables_from_pdf(pdf_path, MORTH_2022_DIR)139        total_tables += n140        print(f"  => {n} tables extracted from {pdf_path.name}")141 142    print(f"\nDone: {total_tables} total table CSVs written to {MORTH_2022_DIR}")143    print("These CSVs can now be used by seed_blackspots.py for accident data seeding.")144 145 146if __name__ == "__main__":147    main()148