Upload caption_integration.py with huggingface_hub

caption_integration.py

@@ -0,0 +1,168 @@
+#!/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}")