caption_integration.py

#!/usr/bin/env python3
"""
Caption Integration Module
Adds Qwen VL captions to template matching results
"""

from typing import Dict, List, Any
import json
from pathlib import Path
from caption_cropped_images import get_captions

def add_captions_to_matches(
    matches: List[Dict],
    captions: Dict[str, str]
) -> List[Dict]:
    """
    Add captions to template matches.
    
    Args:
        matches: List of matched elements from template matching
        captions: Dictionary of crop_id → caption text
    
    Returns:
        Matches with captions added
    """
    
    for match in matches:
        crop_id = match.get('template_id', '')
        caption = captions.get(crop_id, 'unknown')
        match['caption'] = caption
    
    return matches

def add_captions_to_coordinates(
    coordinates: List[Dict],
    captions: Dict[str, str]
) -> List[Dict]:
    """
    Add captions to coordinates output.
    
    Args:
        coordinates: List of detected elements
        captions: Dictionary of crop_id → caption
    
    Returns:
        Coordinates with captions added
    """
    
    for elem in coordinates:
        crop_id = elem.get('template_id') or elem.get('element_id', '')
        caption = captions.get(crop_id, 'unknown')
        elem['caption'] = caption
    
    return coordinates

def create_labeled_output(
    analysis_result: Dict,
    captions: Dict[str, str],
    include_caption: str = "caption"
) -> Dict:
    """
    Create labeled output with captions.
    
    Args:
        analysis_result: Original analysis result
        captions: Captions dictionary
        include_caption: Field name for caption ('caption', 'label', 'description')
    
    Returns:
        Updated analysis result with captions
    """
    
    if 'analysis' in analysis_result and 'elements' in analysis_result['analysis']:
        for elem in analysis_result['analysis']['elements']:
            crop_id = elem.get('template_id') or elem.get('element_id', '')
            caption = captions.get(crop_id, 'unknown')
            elem[include_caption] = caption
    
    if 'exports' in analysis_result:
        analysis_result['exports']['captions'] = captions
    
    return analysis_result

def generate_csv_with_captions(
    coordinates: List[Dict],
    include_caption: bool = True
) -> str:
    """
    Generate CSV with captions.
    
    Args:
        coordinates: List of detected elements with captions
        include_caption: Whether to include caption column
    
    Returns:
        CSV string
    """
    
    import io
    import csv
    
    output = io.StringIO()
    
    # Define columns
    if include_caption:
        columns = [
            'Element_ID', 'Caption', 'X', 'Y', 'X1', 'Y1', 'X2', 'Y2',
            'Width', 'Height', 'Confidence'
        ]
    else:
        columns = [
            'Element_ID', 'X', 'Y', 'X1', 'Y1', 'X2', 'Y2',
            'Width', 'Height', 'Confidence'
        ]
    
    writer = csv.writer(output)
    writer.writerow(columns)
    
    for i, coord in enumerate(coordinates, 1):
        element_id = coord.get('element_id', f"crop_{i:04d}")
        caption = coord.get('caption', 'unknown')
        x = coord.get('x', 0)
        y = coord.get('y', 0)
        x1 = coord.get('x1', 0)
        y1 = coord.get('y1', 0)
        x2 = coord.get('x2', 0)
        y2 = coord.get('y2', 0)
        width = coord.get('width', 0)
        height = coord.get('height', 0)
        confidence = coord.get('confidence', 0)
        
        if include_caption:
            writer.writerow([
                element_id, caption, x, y, x1, y1, x2, y2,
                width, height, f"{confidence:.4f}"
            ])
        else:
            writer.writerow([
                element_id, x, y, x1, y1, x2, y2,
                width, height, f"{confidence:.4f}"
            ])
    
    return output.getvalue()

if __name__ == "__main__":
    # Example usage
    print("[Test] Caption Integration Module")
    
    # Load test captions
    test_captions = {
        "crop_0001": "close button",
        "crop_0002": "search input",
        "crop_0003": "menu icon"
    }
    
    # Test with sample coordinates
    sample_coords = [
        {
            "element_id": "crop_0001",
            "x": 100,
            "y": 50,
            "confidence": 0.95
        }
    ]
    
    # Add captions
    labeled = add_captions_to_coordinates(sample_coords, test_captions)
    print(f"✓ Added captions: {labeled}")