import torch
from typing import List


def detections_to_dict(
    bboxes: list, image_ids: torch.Tensor, image_size: torch.Tensor
) -> List[dict]:
    """Used to turn raw bounding box detections into a dictionary which can be
    serialized for the pycocotools package.

    Args:
        bboxes: A list of bounding boxes to save.
        image_ids: The ids of the images which the boxes correspond to. We need these
            in order to match predictions to ground truth.
        image_size: The boxes are original normalized, so we need to project them to
            image coordinates.

    Returns:
        A list of box dictionaries in the COCO format.
    """

    detections = []
    for image_boxes, image_id in zip(bboxes, image_ids):
        for bbox in image_boxes:
            box = bbox.box
            # XYXY -> XYWH
            box[2:] -= box[:2]
            box *= image_size
            detections.append(
                {
                    "image_id": image_id.int().item(),
                    "category_id": bbox.class_id,
                    "bbox": box.int().tolist(),
                    "score": bbox.confidence,
                }
            )

    return detections