src/bean_vision/export/coco_exporter.py

"""Export predictions to COCO format."""

import json
from pathlib import Path
from typing import Dict, List, Any, Optional
from datetime import datetime

from .polygon_utils import polygon_area, polygon_bbox


class COCOExporter:
    """Export detection results to COCO format."""
    
    def __init__(self, dataset_name: str = "bean_detections"):
        """Initialize COCO exporter.
        
        Args:
            dataset_name: Name of the dataset
        """
        self.dataset_name = dataset_name
        self.coco_data = self._initialize_coco_structure()
        self.annotation_id = 1
        self.image_id = 1
    
    def _initialize_coco_structure(self) -> Dict[str, Any]:
        """Initialize empty COCO structure."""
        return {
            "info": {
                "description": self.dataset_name,
                "version": "1.0",
                "year": datetime.now().year,
                "date_created": datetime.now().isoformat()
            },
            "categories": [
                {"id": 1, "name": "bean", "supercategory": "object"}
            ],
            "images": [],
            "annotations": []
        }
    
    def add_image(
        self,
        image_path: Path,
        width: int,
        height: int,
        image_id: Optional[int] = None
    ) -> int:
        """Add image to COCO dataset.
        
        Args:
            image_path: Path to image file
            width: Image width
            height: Image height
            image_id: Optional specific image ID
            
        Returns:
            Image ID
        """
        if image_id is None:
            image_id = self.image_id
            self.image_id += 1
        
        image_info = {
            "id": image_id,
            "file_name": image_path.name,
            "width": width,
            "height": height
        }
        
        self.coco_data["images"].append(image_info)
        return image_id
    
    def add_predictions(
        self,
        predictions: Dict[str, Any],
        image_id: int
    ) -> int:
        """Add predictions for an image.
        
        Args:
            predictions: Prediction results with polygons
            image_id: ID of the image
            
        Returns:
            Number of annotations added
        """
        num_added = 0
        
        # Handle polygon predictions
        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:
                        # Flatten polygon
                        segmentation = [coord for point in polygon for coord in point]
                        
                        # Calculate bbox and area
                        x1, y1, x2, y2 = polygon_bbox(polygon)
                        area = polygon_area(polygon)
                        
                        annotation = {
                            "id": self.annotation_id,
                            "image_id": image_id,
                            "category_id": 1,  # bean
                            "segmentation": [segmentation],
                            "area": area,
                            "bbox": [x1, y1, x2 - x1, y2 - y1],
                            "iscrowd": 0,
                            "score": float(score)
                        }
                        
                        self.coco_data["annotations"].append(annotation)
                        self.annotation_id += 1
                        num_added += 1
        
        # Handle box-only predictions
        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
                
                annotation = {
                    "id": self.annotation_id,
                    "image_id": image_id,
                    "category_id": 1,
                    "bbox": [x1, y1, x2 - x1, y2 - y1],
                    "area": (x2 - x1) * (y2 - y1),
                    "iscrowd": 0,
                    "score": float(score)
                }
                
                self.coco_data["annotations"].append(annotation)
                self.annotation_id += 1
                num_added += 1
        
        return num_added
    
    def save(self, output_path: Path):
        """Save COCO data to JSON file.
        
        Args:
            output_path: Path to save JSON file
        """
        with open(output_path, 'w') as f:
            json.dump(self.coco_data, f, indent=2)
        print(f"COCO format saved to {output_path}")
    
    def get_statistics(self) -> Dict[str, int]:
        """Get dataset statistics.
        
        Returns:
            Dictionary with counts
        """
        return {
            "num_images": len(self.coco_data["images"]),
            "num_annotations": len(self.coco_data["annotations"]),
            "num_categories": len(self.coco_data["categories"])
        }