feat: add new examples from pocketflow-academy

2025-03-19 10:31:04 -03:00
parent 84720ceebd
commit 557a14f695
129 changed files with 13455 additions and 0 deletions
@@ -0,0 +1,74 @@
+# PocketFlow Tool: PDF Vision
+
+A PocketFlow example project demonstrating PDF processing with OpenAI's Vision API for OCR and text extraction.
+
+## Features
+
+- Convert PDF pages to images while maintaining quality and size limits
+- Extract text from scanned documents using GPT-4 Vision API
+- Support for custom extraction prompts
+- Maintain page order and formatting in extracted text
+- Batch processing of multiple PDFs from a directory
+
+## Installation
+
+1. Clone the repository
+2. Install dependencies:
+   ```bash
+   pip install -r requirements.txt
+   ```
+3. Set your OpenAI API key as an environment variable:
+   ```bash
+   export OPENAI_API_KEY=your_api_key_here
+   ```
+
+## Usage
+
+1. Place your PDF files in the `pdfs` directory
+2. Run the example:
+   ```bash
+   python main.py
+   ```
+   The script will process all PDF files in the `pdfs` directory and output the extracted text for each one.
+
+## Project Structure
+
+```
+pocketflow-tool-pdf-vision/
+├── pdfs/           # Directory for PDF files to process
+├── tools/
+│   ├── pdf.py     # PDF to image conversion
+│   └── vision.py  # Vision API integration
+├── utils/
+│   └── call_llm.py # OpenAI client config
+├── nodes.py       # PocketFlow nodes
+├── flow.py        # Flow configuration
+└── main.py        # Example usage
+```
+
+## Flow Description
+
+1. **LoadPDFNode**: Loads PDF and converts pages to images
+2. **ExtractTextNode**: Processes images with Vision API
+3. **CombineResultsNode**: Combines extracted text from all pages
+
+## Customization
+
+You can customize the extraction by modifying the prompt in `shared`:
+
+```python
+shared = {
+    "pdf_path": "your_file.pdf",
+    "extraction_prompt": "Your custom prompt here"
+}
+```
+
+## Limitations
+
+- Maximum PDF page size: 2000px (configurable in `tools/pdf.py`)
+- Vision API token limit: 1000 tokens per response
+- Image size limit: 20MB per image for Vision API
+
+## License
+
+MIT
@@ -0,0 +1,6 @@
+from pocketflow import Flow
+from nodes import ProcessPDFBatchNode
+
+def create_vision_flow():
+    """Create a flow for batch PDF processing with Vision API"""
+    return Flow(start=ProcessPDFBatchNode())
@@ -0,0 +1,17 @@
+from flow import create_vision_flow
+
+def main():
+    # Create and run flow
+    flow = create_vision_flow()
+    shared = {}
+    flow.run(shared)
+    
+    # Print results
+    if "results" in shared:
+        for result in shared["results"]:
+            print(f"\nFile: {result['filename']}")
+            print("-" * 50)
+            print(result["text"])
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,127 @@
+from pocketflow import Node, BatchNode
+from tools.pdf import pdf_to_images
+from tools.vision import extract_text_from_image
+from typing import List, Dict, Any
+from pathlib import Path
+import os
+
+class ProcessPDFBatchNode(BatchNode):
+    """Node for processing multiple PDFs from a directory"""
+    
+    def prep(self, shared):
+        # Get PDF directory path
+        root_dir = Path(__file__).parent
+        pdf_dir = root_dir / "pdfs"
+        
+        # List all PDFs
+        pdf_files = []
+        for file in os.listdir(pdf_dir):
+            if file.lower().endswith('.pdf'):
+                pdf_files.append({
+                    "pdf_path": str(pdf_dir / file),
+                    "extraction_prompt": shared.get("extraction_prompt", 
+                        "Extract all text from this document, preserving formatting and layout.")
+                })
+        
+        if not pdf_files:
+            print("No PDF files found in 'pdfs' directory!")
+            return []
+            
+        print(f"Found {len(pdf_files)} PDF files")
+        return pdf_files
+    
+    def exec(self, item):
+        # Create flow for single PDF
+        flow = create_single_pdf_flow()
+        
+        # Process PDF
+        print(f"\nProcessing: {os.path.basename(item['pdf_path'])}")
+        print("-" * 50)
+        
+        # Run flow
+        shared = item.copy()
+        flow.run(shared)
+        
+        return {
+            "filename": os.path.basename(item["pdf_path"]),
+            "text": shared.get("final_text", "No text extracted")
+        }
+    
+    def post(self, shared, prep_res, exec_res_list):
+        shared["results"] = exec_res_list
+        return "default"
+
+class LoadPDFNode(Node):
+    """Node for loading and converting a single PDF to images"""
+    
+    def prep(self, shared):
+        return shared.get("pdf_path", "")
+        
+    def exec(self, pdf_path):
+        return pdf_to_images(pdf_path)
+        
+    def post(self, shared, prep_res, exec_res):
+        shared["page_images"] = exec_res
+        return "default"
+
+class ExtractTextNode(Node):
+    """Node for extracting text from images using Vision API"""
+    
+    def prep(self, shared):
+        return (
+            shared.get("page_images", []),
+            shared.get("extraction_prompt", None)
+        )
+        
+    def exec(self, inputs):
+        images, prompt = inputs
+        results = []
+        
+        for img, page_num in images:
+            text = extract_text_from_image(img, prompt)
+            results.append({
+                "page": page_num,
+                "text": text
+            })
+            
+        return results
+        
+    def post(self, shared, prep_res, exec_res):
+        shared["extracted_text"] = exec_res
+        return "default"
+
+class CombineResultsNode(Node):
+    """Node for combining and formatting extracted text"""
+    
+    def prep(self, shared):
+        return shared.get("extracted_text", [])
+        
+    def exec(self, results):
+        # Sort by page number
+        sorted_results = sorted(results, key=lambda x: x["page"])
+        
+        # Combine text with page numbers
+        combined = []
+        for result in sorted_results:
+            combined.append(f"=== Page {result['page']} ===\n{result['text']}\n")
+            
+        return "\n".join(combined)
+        
+    def post(self, shared, prep_res, exec_res):
+        shared["final_text"] = exec_res
+        return "default"
+
+def create_single_pdf_flow():
+    """Create a flow for processing a single PDF"""
+    from pocketflow import Flow
+    
+    # Create nodes
+    load_pdf = LoadPDFNode()
+    extract_text = ExtractTextNode()
+    combine_results = CombineResultsNode()
+    
+    # Connect nodes
+    load_pdf >> extract_text >> combine_results
+    
+    # Create and return flow
+    return Flow(start=load_pdf)
@@ -0,0 +1,4 @@
+pocketflow>=0.1.0
+openai>=1.0.0
+PyMuPDF>=1.22.0  # for PDF processing
+Pillow>=10.0.0   # for image processing
@@ -0,0 +1,52 @@
+import fitz  # PyMuPDF
+from PIL import Image
+import io
+import base64
+from typing import List, Tuple
+
+def pdf_to_images(pdf_path: str, max_size: int = 2000) -> List[Tuple[Image.Image, int]]:
+    """Convert PDF pages to PIL Images with size limit
+    
+    Args:
+        pdf_path (str): Path to PDF file
+        max_size (int): Maximum dimension (width/height) for images
+        
+    Returns:
+        list: List of tuples (PIL Image, page number)
+    """
+    doc = fitz.open(pdf_path)
+    images = []
+    
+    try:
+        for page_num in range(len(doc)):
+            page = doc[page_num]
+            pix = page.get_pixmap()
+            
+            # Convert to PIL Image
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            
+            # Resize if needed while maintaining aspect ratio
+            if max(img.size) > max_size:
+                ratio = max_size / max(img.size)
+                new_size = tuple(int(dim * ratio) for dim in img.size)
+                img = img.resize(new_size, Image.Resampling.LANCZOS)
+            
+            images.append((img, page_num + 1))
+            
+    finally:
+        doc.close()
+        
+    return images
+
+def image_to_base64(image: Image.Image) -> str:
+    """Convert PIL Image to base64 string
+    
+    Args:
+        image (PIL.Image): Image to convert
+        
+    Returns:
+        str: Base64 encoded image string
+    """
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode('utf-8')
@@ -0,0 +1,40 @@
+from PIL import Image
+from utils.call_llm import client
+from tools.pdf import image_to_base64
+
+def extract_text_from_image(image: Image.Image, prompt: str = None) -> str:
+    """Extract text from image using OpenAI Vision API
+    
+    Args:
+        image (PIL.Image): Image to process
+        prompt (str, optional): Custom prompt for extraction. Defaults to general OCR.
+        
+    Returns:
+        str: Extracted text from image
+    """
+    # Convert image to base64
+    img_base64 = image_to_base64(image)
+    
+    # Default prompt for general OCR
+    if prompt is None:
+        prompt = "Please extract all text from this image."
+    
+    # Call Vision API
+    response = client.chat.completions.create(
+        model="gpt-4o",
+        messages=[{
+            "role": "user",
+            "content": [
+                {"type": "text", "text": prompt},
+                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_base64}"}}
+            ]
+        }]
+    )
+    
+    return response.choices[0].message.content
+
+if __name__ == "__main__":
+    # Test vision processing
+    test_image = Image.open("example.png")
+    result = extract_text_from_image(test_image)
+    print("Extracted text:", result)
@@ -0,0 +1,9 @@
+import os
+from openai import OpenAI
+from pathlib import Path
+
+# Get the project root directory (parent of utils directory)
+ROOT_DIR = Path(__file__).parent.parent
+
+# Initialize OpenAI client with API key from environment
+client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))