| |
| """ |
| 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() |
| |
| |
| 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__": |
| |
| print("[Test] Caption Integration Module") |
| |
| |
| test_captions = { |
| "crop_0001": "close button", |
| "crop_0002": "search input", |
| "crop_0003": "menu icon" |
| } |
| |
| |
| sample_coords = [ |
| { |
| "element_id": "crop_0001", |
| "x": 100, |
| "y": 50, |
| "confidence": 0.95 |
| } |
| ] |
| |
| |
| labeled = add_captions_to_coordinates(sample_coords, test_captions) |
| print(f"✓ Added captions: {labeled}") |
|
|
