process_all_pdfs.py

#!/usr/bin/env python3
"""
Process all PDF files in the documents folder and generate a single JSON file
in the required format with annotations for signatures and stamps.
"""

import json
import sys
from pathlib import Path
from typing import Dict, Any, List, Tuple

from pipeline import process_pdf_pipeline

def convert_to_annotation_format(
    pipeline_result: Dict[str, Any],
    annotation_id_start: int = 1
) -> Tuple[Dict[str, Any], int]:
    """
    Convert pipeline result to the required annotation format.
    
    Args:
        pipeline_result: Result from process_pdf_pipeline
        annotation_id_start: Starting annotation ID number
        
    Returns:
        Tuple of (formatted_result, next_annotation_id)
    """
    pdf_filename = pipeline_result["pdf"]
    pages_data = pipeline_result["pages"]
    
    result = {}
    annotation_counter = annotation_id_start
    
    for page_data in pages_data:
        page_num = page_data.get("page_number", 1)
        page_key = f"page_{page_num}"
        
        # Get image dimensions
        img_dims = page_data.get("image_dimensions", {})
        width = img_dims.get("width", 0)
        height = img_dims.get("height", 0)
        
        # Collect all annotations (signatures and stamps only, no QR codes)
        annotations = []
        
        # Process signatures
        signatures = page_data.get("signatures", [])
        for sig in signatures:
            bbox = sig.get("bbox", {})
            if bbox:
                x1 = bbox.get("x1", 0)
                y1 = bbox.get("y1", 0)
                width_bbox = bbox.get("width", 0)
                height_bbox = bbox.get("height", 0)
                
                # Calculate area
                area = width_bbox * height_bbox
                
                annotation = {
                    f"annotation_{annotation_counter}": {
                        "category": "signature",
                        "bbox": {
                            "x": float(x1),
                            "y": float(y1),
                            "width": float(width_bbox),
                            "height": float(height_bbox)
                        },
                        "area": float(area)
                    }
                }
                annotations.append(annotation)
                annotation_counter += 1
        
        # Process stamps
        stamps = page_data.get("stamps", [])
        for stamp in stamps:
            bbox = stamp.get("bbox", {})
            if bbox:
                x1 = bbox.get("x1", 0)
                y1 = bbox.get("y1", 0)
                width_bbox = bbox.get("width", 0)
                height_bbox = bbox.get("height", 0)
                
                # Calculate area
                area = width_bbox * height_bbox
                
                annotation = {
                    f"annotation_{annotation_counter}": {
                        "category": "stamp",
                        "bbox": {
                            "x": float(x1),
                            "y": float(y1),
                            "width": float(width_bbox),
                            "height": float(height_bbox)
                        },
                        "area": float(area)
                    }
                }
                annotations.append(annotation)
                annotation_counter += 1
        
        # Only include pages that have annotations
        if annotations:
            result[page_key] = {
                "annotations": annotations,
                "page_size": {
                    "width": int(width),
                    "height": int(height)
                }
            }
    
    return result, annotation_counter


def process_all_pdfs(
    documents_dir: str = "documents",
    output_file: str = "all_annotations.json",
    stamp_model_path: str = "stamp_detector/stamp_model.pt",
    stamp_conf: float = 0.25,
    dpi: int = 200
) -> None:
    """
    Process all PDF files in the documents folder and generate a single JSON file.
    
    Args:
        documents_dir: Directory containing PDF files
        output_file: Output JSON file path
        stamp_model_path: Path to stamp model
        stamp_conf: Confidence threshold for stamp detection
        dpi: DPI for PDF to image conversion
    """
    documents_path = Path(documents_dir)
    
    if not documents_path.exists():
        print(f"Error: Documents directory '{documents_dir}' not found!")
        sys.exit(1)
    
    # Find all PDF files
    pdf_files = sorted(list(documents_path.glob("*.pdf")))
    
    if not pdf_files:
        print(f"No PDF files found in '{documents_dir}' directory!")
        sys.exit(1)
    
    print(f"Found {len(pdf_files)} PDF file(s) to process\n")
    print("=" * 70)
    
    # Final result dictionary
    final_result = {}
    annotation_counter = 1
    
    # Process each PDF
    for i, pdf_file in enumerate(pdf_files, 1):
        print(f"\n[{i}/{len(pdf_files)}] Processing: {pdf_file.name}")
        print("-" * 70)
        
        try:
            # Process PDF using pipeline
            pipeline_result = process_pdf_pipeline(
                pdf_path=str(pdf_file),
                output_dir="pipeline_outputs",
                stamp_model_path=stamp_model_path,
                stamp_conf=stamp_conf,
                dpi=dpi,
                save_intermediate=False
            )
            
            # Convert to annotation format
            pdf_annotations, annotation_counter = convert_to_annotation_format(
                pipeline_result,
                annotation_id_start=annotation_counter
            )
            
            # Only add to result if there are annotations
            if pdf_annotations:
                final_result[pdf_file.name] = pdf_annotations
                print(f"✓ Processed: {len(pdf_annotations)} page(s) with annotations")
            else:
                print(f"⚠ No annotations found in {pdf_file.name}")
            
        except Exception as e:
            print(f"✗ Error processing {pdf_file.name}: {str(e)}")
            import traceback
            traceback.print_exc()
            continue
    
    # Save to JSON file
    output_path = Path(output_file)
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(final_result, f, indent=2, ensure_ascii=False)
    
    print("\n" + "=" * 70)
    print("PROCESSING COMPLETE")
    print("=" * 70)
    print(f"Total PDFs processed: {len(pdf_files)}")
    print(f"PDFs with annotations: {len(final_result)}")
    print(f"Total annotations: {annotation_counter - 1}")
    print(f"Output saved to: {output_path.absolute()}")
    print("=" * 70)


if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(
        description="Process all PDF files in documents folder and generate annotations JSON"
    )
    parser.add_argument(
        "--documents-dir",
        default="documents",
        help="Directory containing PDF files (default: documents)"
    )
    parser.add_argument(
        "--output",
        default="all_annotations.json",
        help="Output JSON file path (default: all_annotations.json)"
    )
    parser.add_argument(
        "--stamp-model",
        default="stamp_detector/stamp_model.pt",
        help="Path to stamp model (default: stamp_detector/stamp_model.pt)"
    )
    parser.add_argument(
        "--stamp-conf",
        type=float,
        default=0.25,
        help="Confidence threshold for stamp detection (default: 0.25)"
    )
    parser.add_argument(
        "--dpi",
        type=int,
        default=200,
        help="DPI for PDF to image conversion (default: 200)"
    )
    
    args = parser.parse_args()
    
    process_all_pdfs(
        documents_dir=args.documents_dir,
        output_file=args.output,
        stamp_model_path=args.stamp_model,
        stamp_conf=args.stamp_conf,
        dpi=args.dpi
    )