o
    .wi6                     @   s   d dl mZ d dlZd dlmZ esdgZ	 ddejdejdee ded	ejf
d
dZddejde	d	ejfddZ
		 	ddejdejdee dede	d	ejfddZdS )    )OptionalN)_TORCHVISION_AVAILABLE#generalized_intersection_over_unionpredstargetiou_thresholdreplacement_valreturnc                 C   s   | j dks| jd dkrtd| j |j dks |jd dkr(td|j ddlm} |  dkrEtj|jd |jd |jtj	dS | dkr\tj| jd | jd | jtj	dS || |}|d urk||||k < |S )	N      z-Expected preds to be of shape (N, 4) but got z.Expected target to be of shape (N, 4) but got r   )generalized_box_iou)devicedtype)
ndimshape
ValueErrortorchvision.opsr   numeltorchzerosr   float32)r   r   r   r   r   iou r   c/home/ubuntu/sommelier/.venv/lib/python3.10/site-packages/torchmetrics/functional/detection/giou.py_giou_update   s   ""
r   Tr   	aggregatec                 C   s0   |s| S |   dkr|   S tjd| jdS )Nr   g        )r   )r   diagmeanr   tensorr   )r   r   r   r   r   _giou_compute-   s   (r    c                 C   s.   t stdtj dt| |||}t||S )aP	  Compute Generalized Intersection over Union (`GIOU`_) between two sets of boxes.

    Both sets of boxes are expected to be in (x1, y1, x2, y2) format with 0 <= x1 < x2 and 0 <= y1 < y2.

    Args:
        preds:
            The input tensor containing the predicted bounding boxes.
        target:
            The tensor containing the ground truth.
        iou_threshold:
            Optional IoU thresholds for evaluation. If set to `None` the threshold is ignored.
        replacement_val:
            Value to replace values under the threshold with.
        aggregate:
            Return the average value instead of the full matrix of values

    Example::
        By default giou is aggregated across all box pairs e.g. mean along the diagonal of the gIoU matrix:

        >>> import torch
        >>> from torchmetrics.functional.detection import generalized_intersection_over_union
        >>> preds = torch.tensor(
        ...     [
        ...         [296.55, 93.96, 314.97, 152.79],
        ...         [328.94, 97.05, 342.49, 122.98],
        ...         [356.62, 95.47, 372.33, 147.55],
        ...     ]
        ... )
        >>> target = torch.tensor(
        ...     [
        ...         [300.00, 100.00, 315.00, 150.00],
        ...         [330.00, 100.00, 350.00, 125.00],
        ...         [350.00, 100.00, 375.00, 150.00],
        ...     ]
        ... )
        >>> generalized_intersection_over_union(preds, target)
        tensor(0.5638)

    Example::
        By setting `aggregate=False` the full IoU matrix is returned:

        >>> import torch
        >>> from torchmetrics.functional.detection import generalized_intersection_over_union
        >>> preds = torch.tensor(
        ...     [
        ...         [296.55, 93.96, 314.97, 152.79],
        ...         [328.94, 97.05, 342.49, 122.98],
        ...         [356.62, 95.47, 372.33, 147.55],
        ...     ]
        ... )
        >>> target = torch.tensor(
        ...     [
        ...         [300.00, 100.00, 315.00, 150.00],
        ...         [330.00, 100.00, 350.00, 125.00],
        ...         [350.00, 100.00, 375.00, 150.00],
        ...     ]
        ... )
        >>> generalized_intersection_over_union(preds, target, aggregate=False)
        tensor([[ 0.6895, -0.4964, -0.4944],
                [-0.5105,  0.4673, -0.3434],
                [-0.6024, -0.4021,  0.5345]])

    `zf` requires that `torchvision` is installed. Please install with `pip install torchmetrics[detection]`.)r   ModuleNotFoundErrorr   __name__r   r    )r   r   r   r   r   r   r   r   r   r   3   s   F
)r   )T)Nr   T)typingr   r   torchmetrics.utilities.importsr   __doctest_skip__Tensorfloatr   boolr    r   r   r   r   r   <module>   sD   
	