coco.py

"""Handles direct I/O operations for working with COCO-style datasets.

COCO-style format specification:
- JSON annotation files containing images, annotations, and categories
- Image directory structure can vary (flat, categorized, nested, multi-source)
- Keypoint annotations with coordinates and visibility flags
- Bounding box and segmentation annotations (polygon and RLE)
- Support for multiple animal categories with different skeletons
- Visibility encoding: binary (0/1) or ternary (0/1/2)
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np

from sleap_io.model.bbox import PredictedBoundingBox, UserBoundingBox
from sleap_io.model.instance import Instance, Track
from sleap_io.model.labeled_frame import LabeledFrame
from sleap_io.model.labels import Labels
from sleap_io.model.mask import PredictedSegmentationMask, UserSegmentationMask
from sleap_io.model.roi import PredictedROI, UserROI
from sleap_io.model.skeleton import Edge, Node, Skeleton
from sleap_io.model.video import Video


def parse_coco_json(json_path: str | Path) -> dict:
    """Parse COCO annotation JSON file and validate structure.

    Args:
        json_path: Path to the COCO annotation JSON file.

    Returns:
        Parsed COCO annotation dictionary.

    Raises:
        FileNotFoundError: If JSON file doesn't exist.
        ValueError: If JSON structure is invalid.
    """
    json_path = Path(json_path)

    if not json_path.exists():
        raise FileNotFoundError(f"COCO annotation file not found: {json_path}")

    with open(json_path, "r") as f:
        data = json.load(f)

    # Validate required COCO fields
    required_fields = ["images", "annotations", "categories"]
    for field in required_fields:
        if field not in data:
            raise ValueError(f"Missing required COCO field: {field}")

    return data


def create_skeleton_from_category(category: dict) -> Skeleton:
    """Create a Skeleton object from a COCO category definition.

    Args:
        category: COCO category dictionary with keypoints and skeleton.

    Returns:
        Skeleton object corresponding to the category.
    """
    if "keypoints" not in category:
        raise ValueError(f"Category '{category['name']}' has no keypoint definitions")

    # Create nodes from keypoint names
    keypoint_names = category["keypoints"]
    nodes = [Node(name) for name in keypoint_names]

    # Create edges from skeleton connections
    edges = []
    if "skeleton" in category:
        for connection in category["skeleton"]:
            if len(connection) == 2:
                # COCO skeleton uses 1-based indexing
                src_idx, dst_idx = connection[0] - 1, connection[1] - 1
                if 0 <= src_idx < len(nodes) and 0 <= dst_idx < len(nodes):
                    edges.append(Edge(nodes[src_idx], nodes[dst_idx]))

    skeleton_name = category.get("name", "unknown")
    return Skeleton(nodes, edges, name=skeleton_name)


def resolve_image_path(image_filename: str, dataset_root: Path) -> Path:
    """Resolve image file path handling various directory structures.

    Args:
        image_filename: Image filename from COCO annotation.
        dataset_root: Root directory of the dataset.

    Returns:
        Resolved absolute path to the image file.

    Raises:
        FileNotFoundError: If image file cannot be found.
    """
    # Try direct path first
    image_path = dataset_root / image_filename
    if image_path.exists():
        return image_path

    # Try common variations
    common_prefixes = ["images", "imgs", "data/images", ""]

    for prefix in common_prefixes:
        if prefix:
            test_path = dataset_root / prefix / image_filename
        else:
            # Try finding the file anywhere in the dataset
            test_path = None
            for found_path in dataset_root.rglob(Path(image_filename).name):
                if found_path.is_file():
                    test_path = found_path
                    break

        if test_path and test_path.exists():
            return test_path

    raise FileNotFoundError(
        f"Image file not found: {image_filename} (searched in {dataset_root})"
    )


def decode_keypoints(
    keypoints: list[float], num_keypoints: int, skeleton: Skeleton
) -> np.ndarray:
    """Decode COCO keypoint format to numpy array for Instance creation.

    Args:
        keypoints: Flat list of [x1, y1, v1, x2, y2, v2, ...] values.
        num_keypoints: Number of keypoints (for validation).
        skeleton: Skeleton object defining the keypoint structure.

    Returns:
        Numpy array of shape (num_keypoints, 3) with [x, y, visibility] values.
    """
    if len(keypoints) != num_keypoints * 3:
        raise ValueError(
            f"Keypoints length {len(keypoints)} doesn't match expected "
            f"{num_keypoints * 3}"
        )

    if len(skeleton.nodes) != num_keypoints:
        raise ValueError(
            f"Skeleton has {len(skeleton.nodes)} nodes but annotation has "
            f"{num_keypoints} keypoints"
        )

    points = []
    for i in range(num_keypoints):
        x = keypoints[i * 3]
        y = keypoints[i * 3 + 1]
        visibility = keypoints[i * 3 + 2]

        # Handle different visibility encodings
        # 0 = not labeled/not visible, 1 = labeled but not visible,
        # 2 = labeled and visible
        # For binary encoding: 0 = not visible, 1 = visible
        if visibility == 0:
            # Not labeled or not visible - use NaN coordinates
            points.append([np.nan, np.nan, False])
        elif visibility == 1:
            # Labeled but not visible (occluded) OR visible (in binary encoding)
            # For now, treat as visible since we can't distinguish binary vs ternary
            points.append([x, y, True])
        elif visibility == 2:
            # Labeled and visible
            points.append([x, y, True])
        else:
            # Unknown visibility value, default to visible
            points.append([x, y, True])

    return np.array(points, dtype=np.float32)


def _decode_coco_rle(counts: list[int], size: list[int]) -> np.ndarray:
    """Decode COCO RLE segmentation to a binary mask.

    COCO RLE uses column-major (Fortran) order. This function decodes the RLE
    counts and returns a row-major numpy array.

    Args:
        counts: RLE counts (alternating runs of 0s and 1s, starting with 0s).
        size: Mask dimensions as [height, width].

    Returns:
        A 2D boolean numpy array of shape (height, width).
    """
    height, width = size
    total = height * width
    flat = np.zeros(total, dtype=bool)
    pos = 0
    for i, count in enumerate(counts):
        if i % 2 == 1:  # Odd indices are 1-runs
            end = min(pos + count, total)
            flat[pos:end] = True
        pos += count
    # COCO RLE is column-major, reshape accordingly
    return flat.reshape((width, height)).T


def _encode_coco_rle(mask: np.ndarray) -> dict:
    """Encode a binary mask as COCO RLE format.

    COCO RLE uses column-major (Fortran) order.

    Args:
        mask: A 2D boolean or uint8 numpy array of shape (height, width).

    Returns:
        COCO RLE dict with "counts" (list of ints) and "size" [height, width].
    """
    height, width = mask.shape
    # Transpose to column-major order then flatten
    flat = mask.T.ravel().astype(np.uint8)

    if len(flat) == 0:
        return {"counts": [], "size": [height, width]}

    # Find positions where value changes
    diff = np.diff(flat)
    change_indices = np.where(diff != 0)[0] + 1

    # Build run lengths
    positions = np.concatenate([[0], change_indices, [len(flat)]])
    run_lengths = np.diff(positions).tolist()

    # Ensure we start with a 0-run
    if flat[0] == 1:
        run_lengths = [0] + run_lengths

    return {"counts": run_lengths, "size": [height, width]}


def _decode_segmentation(
    segmentation: dict | list | None,
    height: int,
    width: int,
    segmentation_format: str,
    score: float | None = None,
    **kwargs,
) -> tuple[list, list]:
    """Decode a COCO ``segmentation`` field into masks and/or ROIs.

    COCO encodes segmentation either as RLE (a dict with ``counts``/``size``) or
    as polygons (a list of flat ``[x1, y1, x2, y2, ...]`` rings, where multiple
    rings belong to a single object). RLE is always materialized as a
    segmentation mask. Polygon handling depends on ``segmentation_format``:

    - ``"mask"``: rasterize all rings of the annotation into a single
      segmentation mask at the image resolution. Requires positive ``height``
      and ``width``; if either is missing (``<= 0``), the polygon is kept as
      ROI(s) instead since rasterization needs the image extent.
    - ``"roi"``: keep the native vector geometry as one ROI per ring.

    When ``score`` is provided, predicted variants
    (`PredictedSegmentationMask` / `PredictedROI`) carrying the score are
    created; otherwise user variants are created.

    Args:
        segmentation: The COCO ``segmentation`` field (RLE dict, polygon list, or
            ``None``).
        height: Image height in pixels (used to rasterize polygons in mask mode).
        width: Image width in pixels (used to rasterize polygons in mask mode).
        segmentation_format: Either ``"mask"`` or ``"roi"``.
        score: Optional COCO detection confidence. If not ``None``, predicted
            mask/ROI variants are created with this score.
        **kwargs: Metadata forwarded to the created mask/ROI objects (e.g.
            ``category``, ``instance``, ``track``).

    Returns:
        A ``(masks, rois)`` tuple of the decoded objects for this annotation.
    """
    masks = []
    rois = []

    if segmentation is None:
        return masks, rois

    # Predicted variants carry a score; user variants do not.
    predicted = score is not None
    mask_cls = PredictedSegmentationMask if predicted else UserSegmentationMask
    roi_cls = PredictedROI if predicted else UserROI
    score_kw = {"score": score} if predicted else {}

    if isinstance(segmentation, dict):
        # RLE format: always materialize as a segmentation mask.
        mask = _decode_coco_rle(segmentation["counts"], segmentation["size"])
        masks.append(mask_cls.from_numpy(mask, **kwargs, **score_kw))
        return masks, rois

    if isinstance(segmentation, list) and len(segmentation) > 0:
        # Polygon format: each entry is a flat [x1, y1, x2, y2, ...] ring. All
        # rings of one annotation describe a single object. Skip degenerate rings
        # with fewer than 3 vertices (points/lines), which cannot form a polygon.
        polygons = []
        for ring in segmentation:
            coords = [
                (float(ring[i]), float(ring[i + 1])) for i in range(0, len(ring) - 1, 2)
            ]
            if len(coords) >= 3:
                polygons.append(coords)

        if not polygons:
            return masks, rois

        if segmentation_format == "mask" and height > 0 and width > 0:
            if len(polygons) == 1:
                roi = roi_cls.from_polygon(polygons[0], **kwargs, **score_kw)
            else:
                roi = roi_cls.from_multi_polygon(polygons, **kwargs, **score_kw)
            # to_mask preserves the predicted/user variant and score from the ROI.
            masks.append(roi.to_mask(height, width))
        else:
            # ROI mode, or mask mode without image dimensions to rasterize into.
            for coords in polygons:
                rois.append(roi_cls.from_polygon(coords, **kwargs, **score_kw))

    return masks, rois


def read_labels(
    json_path: str | Path,
    dataset_root: str | Path | None = None,
    grayscale: bool = False,
    segmentation_format: str = "mask",
    category_as_track: bool = False,
) -> Labels:
    """Read COCO-style dataset and return a Labels object.

    Supports both pose estimation datasets (with keypoints) and detection-only
    datasets (with bounding boxes and/or segmentation masks). Annotations that
    contain both keypoints and segmentation/bbox data will have both preserved:
    keypoints are stored as `Instance` objects while segmentation and bounding box
    data are stored as `ROI` or `SegmentationMask` objects.

    Args:
        json_path: Path to the COCO annotation JSON file.
        dataset_root: Root directory of the dataset. If None, uses parent directory
                     of json_path.
        grayscale: If True, load images as grayscale (1 channel). If False, load as
                   RGB (3 channels). Default is False.
        segmentation_format: How to represent polygon segmentation. ``"mask"`` (the
            default) rasterizes each annotation's polygon(s) into a single
            `SegmentationMask` at the image resolution; ``"roi"`` keeps the native
            vector geometry as `ROI` objects. RLE segmentation is always read as a
            `SegmentationMask` regardless of this setting. In ``"mask"`` mode,
            polygons for images that lack ``height``/``width`` fall back to ROIs
            since rasterization requires the image extent.
        category_as_track: If True, treat each COCO category as a persistent
            identity: one `Track` (named after the category) is created per
            category and assigned to every annotation of that category (masks,
            ROIs, bounding boxes, and keypoint instances without an explicit
            track id). Useful for instance-segmentation datasets where the
            category encodes identity rather than object class. Default is False.

    Returns:
        Parsed labels as a Labels instance.

    Raises:
        ValueError: If ``segmentation_format`` is not ``"mask"`` or ``"roi"``.
    """
    if segmentation_format not in ("mask", "roi"):
        raise ValueError(
            f"segmentation_format must be 'mask' or 'roi', got {segmentation_format!r}."
        )

    # One shared Track per category name, created on first use.
    category_track_dict: dict[str, Track] = {}

    def _category_track(cat_name: str) -> Track | None:
        """Return the shared Track for a category, creating it on first use."""
        if not category_as_track or not cat_name:
            return None
        if cat_name not in category_track_dict:
            category_track_dict[cat_name] = Track(name=cat_name)
        return category_track_dict[cat_name]

    json_path = Path(json_path)

    if dataset_root is None:
        dataset_root = json_path.parent
    else:
        dataset_root = Path(dataset_root)

    # Parse COCO annotation file
    coco_data = parse_coco_json(json_path)

    # Create skeletons from categories and category name mapping
    skeletons = {}
    category_names = {}
    for category in coco_data["categories"]:
        category_names[category["id"]] = category.get("name", "")
        if "keypoints" in category and len(category["keypoints"]) > 0:
            skeleton = create_skeleton_from_category(category)
            skeletons[category["id"]] = skeleton

    # Track management: maps track_id -> Track object
    track_dict = {}

    # Create image id to annotation mapping
    image_annotations = {}
    for annotation in coco_data["annotations"]:
        image_id = annotation["image_id"]
        if image_id not in image_annotations:
            image_annotations[image_id] = []
        image_annotations[image_id].append(annotation)

    # Group images by shape (height, width) for shared Video objects. Each
    # ``images`` entry becomes its own frame in its shape's video, so the frame
    # index is the entry's position within that group. This is keyed by image_id
    # directly (not by resolved path) so distinct images that share a file_name
    # do not collide.
    shape_to_images = {}
    image_id_to_path = {}
    image_id_to_shape = {}
    image_id_to_frame_idx = {}

    for image_info in coco_data["images"]:
        image_id = image_info["id"]
        image_filename = image_info["file_name"]
        height = image_info.get("height", 0)
        width = image_info.get("width", 0)

        # Resolve image path
        try:
            image_path = resolve_image_path(image_filename, dataset_root)
            image_id_to_path[image_id] = image_path

            # Group by shape
            shape_key = (height, width)
            image_id_to_shape[image_id] = shape_key
            if shape_key not in shape_to_images:
                shape_to_images[shape_key] = []
            image_id_to_frame_idx[image_id] = len(shape_to_images[shape_key])
            shape_to_images[shape_key].append(str(image_path))
        except FileNotFoundError:
            # Skip missing images
            continue

    # Create Video objects for each unique shape
    shape_to_video = {}
    for shape_key, image_paths in shape_to_images.items():
        height, width = shape_key
        # Create Video from the list of images with this shape
        video = Video.from_filename(
            image_paths,
            grayscale=grayscale,
        )
        # Store shape metadata from JSON (useful when images can't be read)
        channels = 1 if grayscale else 3
        video.backend_metadata["shape"] = (len(image_paths), height, width, channels)
        shape_to_video[shape_key] = video

    # Process images and annotations
    labeled_frames = []
    rois = []
    masks = []
    bboxes = []

    for image_info in coco_data["images"]:
        image_id = image_info["id"]

        # Skip if image was not found
        if image_id not in image_id_to_path:
            continue

        # Get the video and frame index for this image
        shape_key = image_id_to_shape[image_id]
        img_height, img_width = shape_key
        video = shape_to_video[shape_key]
        frame_idx = image_id_to_frame_idx[image_id]

        # Create instances from annotations
        instances = []
        if image_id in image_annotations:
            for annotation in image_annotations[image_id]:
                category_id = annotation["category_id"]
                cat_name = category_names.get(category_id, "")
                has_kpts = "keypoints" in annotation and annotation["keypoints"]

                if has_kpts and category_id in skeletons:
                    # Pose annotation with keypoints
                    skeleton = skeletons[category_id]

                    # Extract track ID
                    track = None
                    track_id = (
                        annotation.get("attributes", {}).get("object_id")
                        or annotation.get("track_id")
                        or annotation.get("instance_id")
                    )

                    if track_id is not None:
                        if track_id not in track_dict:
                            track_dict[track_id] = Track(name=f"track_{track_id}")
                        track = track_dict[track_id]
                    else:
                        # Fall back to category identity when requested.
                        track = _category_track(cat_name)

                    keypoints_data = annotation["keypoints"]
                    expected_keypoints = len(skeleton.nodes)

                    points_array = decode_keypoints(
                        keypoints_data, expected_keypoints, skeleton
                    )
                    instance = Instance.from_numpy(
                        points_data=points_array,
                        skeleton=skeleton,
                        track=track,
                    )
                    instances.append(instance)

                    # Also extract segmentation/bbox if present alongside
                    # keypoints. Linked annotations share the instance's track.
                    roi_kwargs = dict(
                        category=cat_name,
                        instance=instance,
                        track=track,
                    )

                    segmentation = annotation.get("segmentation")
                    seg_masks, seg_rois = _decode_segmentation(
                        segmentation,
                        img_height,
                        img_width,
                        segmentation_format,
                        **roi_kwargs,
                    )
                    masks.extend(seg_masks)
                    rois.extend(seg_rois)

                    # Create BoundingBox linked to instance if bbox present
                    bbox = annotation.get("bbox")
                    if bbox is not None:
                        x, y, w, h = bbox
                        bbox_obj = UserBoundingBox.from_xywh(
                            x,
                            y,
                            w,
                            h,
                            category=cat_name,
                            instance=instance,
                            track=track,
                        )
                        bboxes.append(bbox_obj)
                else:
                    # Detection-only annotation: create ROIs/masks/bboxes. A
                    # COCO ``score`` marks a prediction, so it selects predicted
                    # mask/ROI/bbox variants (mirrors the bbox handling below).
                    roi_kwargs = dict(
                        category=cat_name,
                        track=_category_track(cat_name),
                    )

                    # Handle segmentation field
                    segmentation = annotation.get("segmentation")
                    seg_masks, seg_rois = _decode_segmentation(
                        segmentation,
                        img_height,
                        img_width,
                        segmentation_format,
                        score=annotation.get("score"),
                        **roi_kwargs,
                    )
                    masks.extend(seg_masks)
                    rois.extend(seg_rois)

                    # Create BoundingBox if no segmentation was processed
                    bbox = annotation.get("bbox")
                    if bbox is not None and not segmentation:
                        x, y, w, h = bbox
                        # COCO score field is only present in prediction results,
                        # so its presence distinguishes predicted from user
                        # annotations.
                        score = annotation.get("score")
                        if score is not None:
                            bbox_obj = PredictedBoundingBox.from_xywh(
                                x,
                                y,
                                w,
                                h,
                                score=score,
                                **roi_kwargs,
                            )
                        else:
                            bbox_obj = UserBoundingBox.from_xywh(
                                x, y, w, h, **roi_kwargs
                            )
                        bboxes.append(bbox_obj)

        # Always create a labeled frame so unannotated images are preserved
        labeled_frame = LabeledFrame(
            video=video, frame_idx=frame_idx, instances=instances
        )
        labeled_frame.rois.extend(rois)
        labeled_frame.masks.extend(masks)
        labeled_frame.bboxes.extend(bboxes)
        labeled_frames.append(labeled_frame)

        # Reset per-frame annotation lists
        rois = []
        masks = []
        bboxes = []

    return Labels(labeled_frames=labeled_frames)


def read_labels_set(
    dataset_path: str | Path,
    json_files: list[str] | None = None,
    grayscale: bool = False,
    segmentation_format: str = "mask",
    category_as_track: bool = False,
) -> dict[str, Labels]:
    """Read multiple COCO annotation files and return a dictionary of Labels.

    This function is designed to handle datasets with multiple splits (train/val/test)
    or multiple annotation files.

    Args:
        dataset_path: Root directory containing COCO annotation files.
        json_files: List of specific JSON filenames to load. If None, automatically
                   discovers all .json files in the dataset directory.
        grayscale: If True, load images as grayscale (1 channel). If False, load as
                   RGB (3 channels). Default is False.
        segmentation_format: How to represent polygon segmentation. ``"mask"`` (the
            default) rasterizes polygons into `SegmentationMask` objects; ``"roi"``
            keeps them as vector `ROI` objects. RLE segmentation is always read as a
            `SegmentationMask`. See `read_labels` for details.
        category_as_track: If True, treat each COCO category as a persistent
            identity, creating one `Track` per category and assigning it to that
            category's annotations. See `read_labels` for details. Tracks are
            created independently per split. Default is False.

    Returns:
        Dictionary mapping split names to Labels objects.
    """
    dataset_path = Path(dataset_path)

    if json_files is None:
        # Auto-discover JSON files
        json_files = [f.name for f in dataset_path.glob("*.json")]
        if not json_files:
            raise FileNotFoundError(f"No JSON annotation files found in {dataset_path}")

    labels_dict = {}

    for json_file in json_files:
        json_path = dataset_path / json_file

        # Use filename (without extension) as split name
        split_name = json_path.stem

        # Load labels for this split
        labels = read_labels(
            json_path,
            dataset_root=dataset_path,
            grayscale=grayscale,
            segmentation_format=segmentation_format,
            category_as_track=category_as_track,
        )
        labels_dict[split_name] = labels

    return labels_dict


def encode_keypoints(
    points_array: np.ndarray, visibility_encoding: str = "ternary"
) -> list[float]:
    """Encode numpy array of points into COCO keypoint format.

    Args:
        points_array: Numpy array of shape (num_keypoints, 2) or (num_keypoints, 3)
                     with [x, y] or [x, y, visibility] values.
        visibility_encoding: Visibility encoding to use. Either "binary" (0/1) or
                           "ternary" (0/1/2). Default is "ternary".

    Returns:
        Flat list of [x1, y1, v1, x2, y2, v2, ...] values.
    """
    keypoints = []
    for i in range(len(points_array)):
        if points_array.shape[1] == 2:
            x, y = points_array[i]
            visible = not (np.isnan(x) or np.isnan(y))
        else:
            x, y = points_array[i, :2]
            visible = points_array[i, 2] if points_array.shape[1] > 2 else True

        # Handle NaN coordinates
        if np.isnan(x) or np.isnan(y):
            keypoints.extend([0.0, 0.0, 0])  # Not labeled
        else:
            # Encode visibility
            if visibility_encoding == "binary":
                # Binary: 0 = not visible, 1 = visible
                visibility_value = 1 if visible else 0
            else:
                # Ternary: 0 = not labeled, 1 = labeled but occluded,
                # 2 = labeled and visible
                visibility_value = 2 if visible else 1
            keypoints.extend([float(x), float(y), visibility_value])

    return keypoints


def convert_labels(
    labels: Labels,
    image_filenames: str | list[str] | None = None,
    visibility_encoding: str = "ternary",
) -> dict:
    """Convert a Labels object into COCO-formatted annotations.

    Args:
        labels: SLEAP `Labels` object to be converted to COCO format.
        image_filenames: Optional image filenames to use. If provided, must be a single
                        string (for single-frame videos) or a list of strings matching
                        the number of labeled frames. If None, generates filenames from
                        video filenames and frame indices.
        visibility_encoding: Visibility encoding to use. Either "binary" (0/1) or
                           "ternary" (0/1/2). Default is "ternary".

    Returns:
        COCO annotation dictionary with "images", "annotations", and "categories"
        fields.

    Note:
        Rotated bounding boxes are not supported by the COCO format. Exporting
        ``BoundingBox`` objects that have a non-zero ``angle`` will raise
        ``ValueError`` when ``BoundingBox.xywh`` is accessed.
    """
    coco_data = {
        "images": [],
        "annotations": [],
        "categories": [],
    }

    # Build skeleton/category mapping
    skeleton_to_category = {}
    category_name_to_id = {}
    category_id_counter = 1
    for skeleton in labels.skeletons:
        if skeleton not in skeleton_to_category:
            cat_name = (
                skeleton.name if skeleton.name else f"skeleton_{category_id_counter}"
            )
            category = {
                "id": category_id_counter,
                "name": cat_name,
                "keypoints": [node.name for node in skeleton.nodes],
                "skeleton": [
                    [i + 1, j + 1]
                    for i, j in [
                        (
                            skeleton.nodes.index(edge.source),
                            skeleton.nodes.index(edge.destination),
                        )
                        for edge in skeleton.edges
                    ]
                ],  # Convert to 1-based indexing
            }
            coco_data["categories"].append(category)
            skeleton_to_category[skeleton] = category_id_counter
            category_name_to_id[cat_name] = category_id_counter
            category_id_counter += 1

    # Build track mapping
    track_to_id = {}
    track_id_counter = 1
    for track in labels.tracks:
        if track not in track_to_id:
            track_to_id[track] = track_id_counter
            track_id_counter += 1

    # Process image filenames
    if image_filenames is not None:
        if isinstance(image_filenames, str):
            image_filenames = [image_filenames]
        if len(image_filenames) != len(labels.labeled_frames):
            raise ValueError(
                f"Number of image filenames ({len(image_filenames)}) must match "
                f"number of labeled frames ({len(labels.labeled_frames)})"
            )

    # Process labeled frames
    image_id_counter = 1
    annotation_id_counter = 1

    # Build mapping from (video, frame_idx) to image_id for ROI/mask export
    video_frame_to_image_id = {}

    for frame_idx, labeled_frame in enumerate(labels.labeled_frames):
        # Determine image filename
        if image_filenames is not None:
            image_filename = image_filenames[frame_idx]
        else:
            # Generate from video filename and frame index
            video = labeled_frame.video
            if isinstance(video.filename, list):
                # Image sequence - use the specific frame
                if labeled_frame.frame_idx < len(video.filename):
                    image_filename = Path(video.filename[labeled_frame.frame_idx]).name
                else:
                    image_filename = f"frame_{labeled_frame.frame_idx:06d}.png"
            else:
                # Video file - generate image name
                video_name = Path(video.filename).stem
                image_filename = f"{video_name}_frame_{labeled_frame.frame_idx:06d}.png"

        # Get image dimensions
        if labeled_frame.video.shape is not None:
            height = labeled_frame.video.shape[1]
            width = labeled_frame.video.shape[2]
        else:
            # Default dimensions if shape unavailable
            height = 0
            width = 0

        # Add image entry
        image_info = {
            "id": image_id_counter,
            "file_name": image_filename,
            "width": width,
            "height": height,
        }
        coco_data["images"].append(image_info)

        # Track video/frame to image_id mapping
        vf_key = (id(labeled_frame.video), labeled_frame.frame_idx)
        video_frame_to_image_id[vf_key] = image_id_counter

        # Process instances
        for instance in labeled_frame.instances:
            # Get category ID
            category_id = skeleton_to_category[instance.skeleton]

            # Encode keypoints
            points_array = instance.numpy()
            keypoints = encode_keypoints(points_array, visibility_encoding)

            # Count visible keypoints
            num_keypoints = sum(
                1 for i in range(0, len(keypoints), 3) if keypoints[i + 2] > 0
            )

            # Compute bounding box from visible keypoints
            visible_points = []
            for i in range(0, len(keypoints), 3):
                if keypoints[i + 2] > 0:  # visible
                    visible_points.append([keypoints[i], keypoints[i + 1]])

            if visible_points:
                visible_points_array = np.array(visible_points)
                x_min = float(np.min(visible_points_array[:, 0]))
                y_min = float(np.min(visible_points_array[:, 1]))
                x_max = float(np.max(visible_points_array[:, 0]))
                y_max = float(np.max(visible_points_array[:, 1]))

                # Bbox in COCO format: [x, y, width, height]
                bbox = [x_min, y_min, x_max - x_min, y_max - y_min]
                area = (x_max - x_min) * (y_max - y_min)
            else:
                # No visible keypoints - use zero bbox
                bbox = [0.0, 0.0, 0.0, 0.0]
                area = 0.0

            # Create annotation
            annotation = {
                "id": annotation_id_counter,
                "image_id": image_id_counter,
                "category_id": category_id,
                "keypoints": keypoints,
                "num_keypoints": num_keypoints,
                "bbox": bbox,
                "area": area,
                "iscrowd": 0,
            }

            # Add track ID if present
            if instance.track is not None:
                annotation["attributes"] = {"object_id": track_to_id[instance.track]}

            coco_data["annotations"].append(annotation)
            annotation_id_counter += 1

        image_id_counter += 1

    # Export ROIs, masks, and bboxes as COCO annotations (iterate by frame)
    for lf in labels.labeled_frames:
        for roi in lf.rois:
            # Get or create category
            cat_name = roi.category if roi.category else "object"
            if cat_name not in category_name_to_id:
                category = {
                    "id": category_id_counter,
                    "name": cat_name,
                }
                coco_data["categories"].append(category)
                category_name_to_id[cat_name] = category_id_counter
                category_id_counter += 1
            category_id = category_name_to_id[cat_name]

            # Find image_id for this ROI
            image_id = _get_or_create_image_id(
                lf.video,
                lf.frame_idx,
                video_frame_to_image_id,
                coco_data,
                image_id_counter,
            )
            if image_id >= image_id_counter:
                image_id_counter = image_id + 1

            annotation = {
                "id": annotation_id_counter,
                "image_id": image_id,
                "category_id": category_id,
                "iscrowd": 0,
            }

            # Write ROI as polygon segmentation with bounding box
            coords = list(roi.geometry.exterior.coords)
            flat = []
            for x, y in coords[:-1]:  # Exclude closing vertex
                flat.extend([float(x), float(y)])
            annotation["segmentation"] = [flat]
            minx, miny, maxx, maxy = roi.bounds
            annotation["bbox"] = [
                minx,
                miny,
                maxx - minx,
                maxy - miny,
            ]
            annotation["area"] = float(roi.area)

            coco_data["annotations"].append(annotation)
            annotation_id_counter += 1

        # Export masks as COCO RLE annotations
        for seg_mask in lf.masks:
            cat_name = seg_mask.category if seg_mask.category else "object"
            if cat_name not in category_name_to_id:
                category = {
                    "id": category_id_counter,
                    "name": cat_name,
                }
                coco_data["categories"].append(category)
                category_name_to_id[cat_name] = category_id_counter
                category_id_counter += 1
            category_id = category_name_to_id[cat_name]

            image_id = _get_or_create_image_id(
                lf.video,
                lf.frame_idx,
                video_frame_to_image_id,
                coco_data,
                image_id_counter,
            )
            if image_id >= image_id_counter:
                image_id_counter = image_id + 1

            # COCO expects full-frame masks; resample if scaled/offset
            export_mask = seg_mask
            if seg_mask.has_spatial_transform:
                target_h, target_w = seg_mask.image_extent
                export_mask = seg_mask.resampled(target_h, target_w)

            mask_data = export_mask.data
            rle = _encode_coco_rle(mask_data)

            bbox_xywh = export_mask.bbox
            annotation = {
                "id": annotation_id_counter,
                "image_id": image_id,
                "category_id": category_id,
                "segmentation": rle,
                "bbox": list(bbox_xywh),
                "area": float(export_mask.area),
                "iscrowd": 1,
            }

            coco_data["annotations"].append(annotation)
            annotation_id_counter += 1

        # Export bounding boxes as COCO bbox annotations (skip instance-linked bboxes
        # since those are already represented in the keypoint annotations above).
        # Note: Rotated bboxes are not supported by COCO format.
        # BoundingBox.xywh raises ValueError for rotated boxes.
        for bbox_obj in lf.bboxes:
            if bbox_obj.instance is not None:
                continue
            cat_name = bbox_obj.category if bbox_obj.category else "object"
            if cat_name not in category_name_to_id:
                category = {
                    "id": category_id_counter,
                    "name": cat_name,
                }
                coco_data["categories"].append(category)
                category_name_to_id[cat_name] = category_id_counter
                category_id_counter += 1
            category_id = category_name_to_id[cat_name]

            image_id = _get_or_create_image_id(
                lf.video,
                lf.frame_idx,
                video_frame_to_image_id,
                coco_data,
                image_id_counter,
            )
            if image_id >= image_id_counter:
                image_id_counter = image_id + 1

            x, y, w, h = bbox_obj.xywh
            annotation = {
                "id": annotation_id_counter,
                "image_id": image_id,
                "category_id": category_id,
                "bbox": [float(x), float(y), float(w), float(h)],
                "area": float(bbox_obj.area),
                "iscrowd": 0,
            }

            if isinstance(bbox_obj, PredictedBoundingBox):
                annotation["score"] = float(bbox_obj.score)

            coco_data["annotations"].append(annotation)
            annotation_id_counter += 1

    return coco_data


def _get_or_create_image_id(
    video: Video | None,
    frame_idx: int | None,
    video_frame_to_image_id: dict,
    coco_data: dict,
    next_image_id: int,
) -> int:
    """Get the image ID for a video/frame pair, creating one if needed.

    Args:
        video: Video object.
        frame_idx: Frame index.
        video_frame_to_image_id: Mapping of (video_id, frame_idx) to image IDs.
        coco_data: COCO data dict to add image entries to.
        next_image_id: Next available image ID.

    Returns:
        The image ID for this video/frame pair.
    """
    if video is not None and frame_idx is not None:
        vf_key = (id(video), frame_idx)
        if vf_key in video_frame_to_image_id:
            return video_frame_to_image_id[vf_key]

    # Create a new image entry
    image_info = {
        "id": next_image_id,
        "file_name": f"frame_{frame_idx if frame_idx is not None else 0:06d}.png",
        "width": 0,
        "height": 0,
    }
    coco_data["images"].append(image_info)

    if video is not None and frame_idx is not None:
        vf_key = (id(video), frame_idx)
        video_frame_to_image_id[vf_key] = next_image_id

    return next_image_id


def write_labels(
    labels: Labels,
    json_path: str | Path,
    image_filenames: str | list[str] | None = None,
    visibility_encoding: str = "ternary",
) -> None:
    """Write Labels to COCO-style JSON annotation file.

    Args:
        labels: SLEAP `Labels` object to save.
        json_path: Path to save the COCO annotation JSON file.
        image_filenames: Optional image filenames to use in the COCO JSON. If
                        provided, must be a single string (for single-frame videos) or
                        a list of strings matching the number of labeled frames. If
                        None, generates filenames from video filenames and frame
                        indices.
        visibility_encoding: Visibility encoding to use. Either "binary" (0/1) or
                           "ternary" (0/1/2). Default is "ternary".

    Notes:
        - This function only writes the JSON annotation file. It does not save images.
        - The generated JSON can be used with mmpose and other COCO-compatible tools.
        - For complete datasets with images, consider using save_dataset() instead.
    """
    json_path = Path(json_path)

    # Convert labels to COCO format
    coco_data = convert_labels(labels, image_filenames, visibility_encoding)

    # Write to JSON file
    json_path.parent.mkdir(parents=True, exist_ok=True)
    with open(json_path, "w") as f:
        json.dump(coco_data, f, indent=2)


def read_coco_panoptic(
    json_path: str | Path,
    images_dir: str | Path | None = None,
) -> Labels:
    """Read COCO panoptic segmentation format.

    Reads the panoptic annotation JSON and per-frame PNG label images.
    Each segment_info entry becomes a LabelImage.Info with:
    - category from the COCO categories table
    - track from the segment id (isthing=True) or None (isthing=False)

    Args:
        json_path: Path to the panoptic annotation JSON.
        images_dir: Directory containing the panoptic PNG files. If None,
            inferred from the JSON path (same directory).

    Returns:
        Labels object with label_images populated.
    """
    from PIL import Image

    from sleap_io.model.label_image import LabelImage, UserLabelImage

    json_path = Path(json_path)
    if images_dir is None:
        images_dir = json_path.parent
    else:
        images_dir = Path(images_dir)

    with open(json_path, "r") as f:
        data = json.load(f)

    # Build category lookup
    categories = {cat["id"]: cat for cat in data.get("categories", [])}

    # Track pool: shared across frames for thing segments
    track_pool: dict[int, Track] = {}

    # Build image_id -> file_name mapping
    image_filenames = []
    image_id_to_idx = {}
    for img in data.get("images", []):
        image_id_to_idx[img["id"]] = len(image_filenames)
        image_filenames.append(img.get("file_name", ""))

    # Create a single video from the image filenames
    video = Video(filename=image_filenames) if image_filenames else Video(filename="")

    labeled_frames = []
    frame_idx = 0
    for ann in data.get("annotations", []):
        png_filename = ann["file_name"]
        segments_info = ann.get("segments_info", [])

        # Read the panoptic PNG and decode to integer label image
        png_path = images_dir / png_filename
        if not png_path.exists():
            continue

        pil_img = Image.open(png_path).convert("RGB")
        rgb = np.array(pil_img, dtype=np.int32)
        # COCO panoptic encoding: pixel_value = R + G * 256 + B * 256^2
        label_data = rgb[:, :, 0] + rgb[:, :, 1] * 256 + rgb[:, :, 2] * 65536

        # Build objects dict from segments_info
        objects: dict[int, LabelImage.Info] = {}
        for seg in segments_info:
            seg_id = seg["id"]
            cat_id = seg["category_id"]
            cat = categories.get(cat_id, {})
            cat_name = cat.get("name", "")
            is_thing = bool(cat.get("isthing", 0))

            track = None
            if is_thing:
                if seg_id not in track_pool:
                    track_pool[seg_id] = Track(name=str(seg_id))
                track = track_pool[seg_id]

            objects[seg_id] = LabelImage.Info(
                track=track,
                category=cat_name,
            )

        li = UserLabelImage(
            data=label_data,
            objects=objects,
        )
        lf = LabeledFrame(video=video, frame_idx=frame_idx)
        lf.label_images.append(li)
        labeled_frames.append(lf)
        frame_idx += 1

    return Labels(labeled_frames=labeled_frames)


def write_coco_panoptic(
    path: str | Path,
    labels: Labels,
    images_dir: str | Path | None = None,
) -> None:
    """Write COCO panoptic segmentation format.

    Writes a panoptic JSON and per-frame PNG label images.

    Args:
        path: Path to save the panoptic annotation JSON.
        labels: Labels object with label_images populated.
        images_dir: Directory to write panoptic PNG files. If None,
            creates a subdirectory next to the JSON named
            ``<json_stem>_panoptic``.
    """
    from PIL import Image

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    if images_dir is None:
        images_dir = path.parent / f"{path.stem}_panoptic"
    else:
        images_dir = Path(images_dir)
    images_dir.mkdir(parents=True, exist_ok=True)

    # Collect all unique categories across all label images
    category_name_to_id: dict[str, int] = {}
    # Track which categories are "thing" (have tracks) vs "stuff" (no track)
    category_is_thing: dict[str, bool] = {}
    cat_id_counter = 1

    for li in labels.label_images:
        for info in li.objects.values():
            cat_name = info.category if info.category else "unknown"
            if cat_name not in category_name_to_id:
                category_name_to_id[cat_name] = cat_id_counter
                # A category is "thing" if any object with it has a track
                category_is_thing[cat_name] = info.track is not None
                cat_id_counter += 1
            else:
                # If any object with this category has a track, it's a thing
                if info.track is not None:
                    category_is_thing[cat_name] = True

    # Build categories list
    coco_categories = []
    for cat_name, cat_id in category_name_to_id.items():
        coco_categories.append(
            {
                "id": cat_id,
                "name": cat_name,
                "isthing": 1 if category_is_thing.get(cat_name, False) else 0,
            }
        )

    # Build images and annotations
    coco_images = []
    coco_annotations = []

    for idx, li in enumerate(labels.label_images):
        image_id = idx + 1
        png_filename = f"panoptic_{image_id:06d}.png"

        # Image entry
        coco_images.append(
            {
                "id": image_id,
                "file_name": f"image_{image_id:06d}.jpg",
                "width": li.width,
                "height": li.height,
            }
        )

        # Encode label data as RGB PNG
        # R = id % 256, G = (id // 256) % 256, B = (id // 65536) % 256
        rgb = np.zeros((li.height, li.width, 3), dtype=np.uint8)
        rgb[:, :, 0] = (li.data % 256).astype(np.uint8)
        rgb[:, :, 1] = ((li.data // 256) % 256).astype(np.uint8)
        rgb[:, :, 2] = ((li.data // 65536) % 256).astype(np.uint8)

        pil_img = Image.fromarray(rgb)
        pil_img.save(images_dir / png_filename)

        # Build segments_info
        segments_info = []
        for seg_id, info in li.objects.items():
            cat_name = info.category if info.category else "unknown"
            cat_id = category_name_to_id[cat_name]
            seg_mask = li.data == seg_id
            area = int(np.sum(seg_mask))

            # Compute bounding box [x, y, width, height] per COCO spec
            ys, xs = np.where(seg_mask)
            if len(xs) > 0:
                bbox = [
                    int(xs.min()),
                    int(ys.min()),
                    int(xs.max() - xs.min()) + 1,
                    int(ys.max() - ys.min()) + 1,
                ]
            else:
                bbox = [0, 0, 0, 0]

            segments_info.append(
                {
                    "id": seg_id,
                    "category_id": cat_id,
                    "area": area,
                    "bbox": bbox,
                    "iscrowd": 0,
                }
            )

        coco_annotations.append(
            {
                "image_id": image_id,
                "file_name": png_filename,
                "segments_info": segments_info,
            }
        )

    coco_data = {
        "images": coco_images,
        "annotations": coco_annotations,
        "categories": coco_categories,
    }

    with open(path, "w") as f:
        json.dump(coco_data, f, indent=2)