models.py

from dataclasses import dataclass
from typing import List, Optional, Tuple

import numpy as np


@dataclass
class Annotation:
    label: str
    score: float
    bbox: List[float]
    mask: Optional[np.ndarray] = None


@dataclass
class ProcessedImage:
    filename: str
    image_b64: str
    annotations: List[Annotation]
    is_annotated: bool
    image_size: Tuple[int, int]  # (width, height)

    def to_export_dict(self) -> dict:
        """Convert to a dict for export, keeping masks for downstream processing."""
        return {
            "filename": self.filename,
            "is_annotated": self.is_annotated,
            "image_size": self.image_size,
            "annotations": [
                {
                    "label": ann.label,
                    "score": ann.score,
                    "bbox": ann.bbox,
                    "mask": ann.mask,
                }
                for ann in self.annotations
            ],
        }