Pre_Seg_Server/backend/routers/export.py

"""Annotation export endpoints (COCO, PNG masks)."""

import io
import json
import logging
import os
import zipfile
from datetime import datetime
from typing import Any, Dict, List

import numpy as np
from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.responses import StreamingResponse
from sqlalchemy.orm import Session

from database import get_db
from models import Project, Annotation, Frame, Template

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/api/export", tags=["Export"])


def _mask_from_polygon(
    polygon: List[List[float]],
    width: int,
    height: int,
) -> np.ndarray:
    """Render a normalized polygon to a binary mask."""
    import cv2

    pts = np.array(
        [[int(p[0] * width), int(p[1] * height)] for p in polygon],
        dtype=np.int32,
    )
    mask = np.zeros((height, width), dtype=np.uint8)
    cv2.fillPoly(mask, [pts], 255)
    return mask


@router.get(
    "/{project_id}/coco",
    summary="Export annotations in COCO format",
)
def export_coco(project_id: int, db: Session = Depends(get_db)) -> StreamingResponse:
    """Export all annotations for a project as a COCO-format JSON file."""
    project = db.query(Project).filter(Project.id == project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    annotations = (
        db.query(Annotation)
        .filter(Annotation.project_id == project_id)
        .all()
    )
    frames = (
        db.query(Frame)
        .filter(Frame.project_id == project_id)
        .order_by(Frame.frame_index)
        .all()
    )
    templates = db.query(Template).all()

    # Build COCO structure
    images = []
    for idx, frame in enumerate(frames):
        images.append({
            "id": frame.id,
            "file_name": frame.image_url,
            "width": frame.width or 1920,
            "height": frame.height or 1080,
            "frame_index": idx,
        })

    categories = []
    template_id_to_cat_id: Dict[int, int] = {}
    for cat_idx, tmpl in enumerate(templates, start=1):
        categories.append({
            "id": cat_idx,
            "name": tmpl.name,
            "color": tmpl.color,
        })
        template_id_to_cat_id[tmpl.id] = cat_idx

    coco_annotations = []
    ann_id = 1
    for ann in annotations:
        if not ann.mask_data:
            continue
        polygons = ann.mask_data.get("polygons", [])
        if not polygons:
            continue

        # Use first polygon for bbox / area approximation
        first_poly = polygons[0]
        xs = [p[0] for p in first_poly]
        ys = [p[1] for p in first_poly]
        width = ann.frame.width if ann.frame else 1920
        height = ann.frame.height if ann.frame else 1080
        bbox = [
            min(xs) * width,
            min(ys) * height,
            (max(xs) - min(xs)) * width,
            (max(ys) - min(ys)) * height,
        ]
        area = bbox[2] * bbox[3]

        segmentation = []
        for poly in polygons:
            flat = []
            for p in poly:
                flat.append(p[0] * width)
                flat.append(p[1] * height)
            segmentation.append(flat)

        coco_annotations.append({
            "id": ann_id,
            "image_id": ann.frame_id,
            "category_id": template_id_to_cat_id.get(ann.template_id, 0),
            "segmentation": segmentation,
            "area": area,
            "bbox": bbox,
            "iscrowd": 0,
        })
        ann_id += 1

    coco = {
        "info": {
            "description": f"Annotations for {project.name}",
            "version": "1.0",
            "year": datetime.now().year,
            "date_created": datetime.now().isoformat(),
        },
        "images": images,
        "annotations": coco_annotations,
        "categories": categories,
    }

    data = json.dumps(coco, ensure_ascii=False, indent=2).encode("utf-8")
    filename = f"project_{project_id}_coco.json"

    return StreamingResponse(
        io.BytesIO(data),
        media_type="application/json",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )


@router.get(
    "/{project_id}/masks",
    summary="Export PNG masks as a ZIP archive",
)
def export_masks(project_id: int, db: Session = Depends(get_db)) -> StreamingResponse:
    """Export all annotation masks as individual PNG files inside a ZIP archive."""
    project = db.query(Project).filter(Project.id == project_id).first()
    if not project:
        raise HTTPException(status_code=404, detail="Project not found")

    annotations = (
        db.query(Annotation)
        .filter(Annotation.project_id == project_id)
        .all()
    )

    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, "w", zipfile.ZIP_DEFLATED) as zf:
        for ann in annotations:
            if not ann.mask_data:
                continue
            polygons = ann.mask_data.get("polygons", [])
            if not polygons:
                continue

            width = ann.frame.width if ann.frame else 1920
            height = ann.frame.height if ann.frame else 1080
            combined = np.zeros((height, width), dtype=np.uint8)

            for poly in polygons:
                mask = _mask_from_polygon(poly, width, height)
                combined = np.maximum(combined, mask)

            # Encode PNG
            import cv2
            _, encoded = cv2.imencode(".png", combined)
            fname = f"mask_{ann.id:06d}.png"
            zf.writestr(fname, encoded.tobytes())

    zip_buffer.seek(0)
    filename = f"project_{project_id}_masks.zip"

    return StreamingResponse(
        zip_buffer,
        media_type="application/zip",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )