src/bean_vision/export/labelme_exporter.py

"""Export predictions to LabelMe format."""

import json
import base64
from pathlib import Path
from typing import Dict, List, Any, Optional
from PIL import Image
import io


class LabelMeExporter:
    """Export detection results to LabelMe format."""
    
    def __init__(self, version: str = "5.0.1"):
        """Initialize LabelMe exporter.
        
        Args:
            version: LabelMe version string
        """
        self.version = version
    
    def create_labelme_annotation(
        self,
        image_path: Path,
        predictions: Dict[str, Any],
        include_image_data: bool = False
    ) -> Dict[str, Any]:
        """Create LabelMe format annotation.
        
        Args:
            image_path: Path to image file
            predictions: Prediction results with polygons
            include_image_data: Whether to include base64 image data
            
        Returns:
            LabelMe format dictionary
        """
        # Load image to get dimensions
        with Image.open(image_path) as img:
            width, height = img.size
            
            # Optionally encode image data
            image_data = None
            if include_image_data:
                buffer = io.BytesIO()
                img.save(buffer, format=img.format or 'PNG')
                image_data = base64.b64encode(buffer.getvalue()).decode('utf-8')
        
        # Create LabelMe structure
        labelme_data = {
            "version": self.version,
            "flags": {},
            "shapes": [],
            "imagePath": str(image_path.name),
            "imageData": image_data,
            "imageHeight": height,
            "imageWidth": width
        }
        
        # Add polygon shapes
        if 'polygons' in predictions:
            polygons = predictions['polygons']
            scores = predictions.get('scores', [1.0] * len(polygons))
            
            for polygon_list, score in zip(polygons, scores):
                for polygon in polygon_list:
                    if len(polygon) >= 3:
                        shape = {
                            "label": "bean",
                            "points": polygon,
                            "group_id": None,
                            "shape_type": "polygon",
                            "flags": {},
                            "attributes": {
                                "confidence": float(score)
                            }
                        }
                        labelme_data["shapes"].append(shape)
        
        # Add boxes as rectangles if no polygons
        elif 'boxes' in predictions:
            boxes = predictions['boxes']
            scores = predictions.get('scores', [1.0] * len(boxes))
            
            for box, score in zip(boxes, scores):
                x1, y1, x2, y2 = box
                
                shape = {
                    "label": "bean",
                    "points": [[x1, y1], [x2, y2]],
                    "group_id": None,
                    "shape_type": "rectangle",
                    "flags": {},
                    "attributes": {
                        "confidence": float(score)
                    }
                }
                labelme_data["shapes"].append(shape)
        
        return labelme_data
    
    def save(
        self,
        image_path: Path,
        predictions: Dict[str, Any],
        output_path: Path,
        include_image_data: bool = False
    ):
        """Save LabelMe annotation to JSON file.
        
        Args:
            image_path: Path to image file
            predictions: Prediction results
            output_path: Path to save JSON file
            include_image_data: Whether to include base64 image data
        """
        labelme_data = self.create_labelme_annotation(
            image_path,
            predictions,
            include_image_data
        )
        
        with open(output_path, 'w') as f:
            json.dump(labelme_data, f, indent=2)
        
        print(f"LabelMe format saved to {output_path}")
    
    def save_batch(
        self,
        results: List[Dict[str, Any]],
        output_dir: Path,
        include_image_data: bool = False
    ):
        """Save batch of predictions to LabelMe format.
        
        Args:
            results: List of prediction results with image_path
            output_dir: Directory to save annotations
            include_image_data: Whether to include base64 image data
        """
        output_dir.mkdir(parents=True, exist_ok=True)
        
        for result in results:
            if 'error' in result:
                continue
            
            image_path = Path(result['image_path'])
            output_path = output_dir / f"{image_path.stem}.json"
            
            self.save(
                image_path,
                result,
                output_path,
                include_image_data
            )
        
        print(f"Saved {len(results)} LabelMe annotations to {output_dir}")