a
    d"                     @  s   d dl mZ d dlmZ d dlmZmZmZmZm	Z	m
Z
 d dlZd dlmZ ddlmZmZ G dd	 d	eZG d
d deZdS )    )annotations)Enum)AnyListOptionalSequenceTupleUnionN)InterpolationMode   )_FillTypeJIT	Datapointc                   @  s   e Zd ZdZdZdZdZdS )BoundingBoxFormatz{[BETA] Coordinate format of a bounding box.

    Available formats are

    * ``XYXY``
    * ``XYWH``
    * ``CXCYWH``
    XYXYXYWHCXCYWHN)__name__
__module____qualname____doc__r   r   r    r   r   m/var/www/html/stable-diffusion-webui/venv/lib/python3.9/site-packages/torchvision/datapoints/_bounding_box.pyr      s   	r   c                
   @  s  e Zd ZU dZded< ded< edddd ddd	Zd
d
d
dddddddd dddZed
d
dd dddd dddZd
ddddddZ	d dd d!Z
d dd"d#Zejd
d$fd%d&d'd(d d)d*d+Zd,d,d,d,d d-d.d/Zd%d d0d1d2Zejd$fd,d,d,d,d%d&d(d d3d4d5ZdPd7d8dd d9d:d;Zejd<d
d
fd=d&d>d?d@d dAdBdCZejd
d
fdDdEd=dEd&d@d?d dFdGdHZejd
d
fdIdId&d@d?d dJdKdLZejd
fdd&d@d dMdNdOZd
S )QBoundingBoxa  [BETA] :class:`torch.Tensor` subclass for bounding boxes.

    Args:
        data: Any data that can be turned into a tensor with :func:`torch.as_tensor`.
        format (BoundingBoxFormat, str): Format of the bounding box.
        spatial_size (two-tuple of ints): Height and width of the corresponding image or video.
        dtype (torch.dtype, optional): Desired data type of the bounding box. If omitted, will be inferred from
            ``data``.
        device (torch.device, optional): Desired device of the bounding box. If omitted and ``data`` is a
            :class:`torch.Tensor`, the device is taken from it. Otherwise, the bounding box is constructed on the CPU.
        requires_grad (bool, optional): Whether autograd should record operations on the bounding box. If omitted and
            ``data`` is a :class:`torch.Tensor`, the value is taken from it. Otherwise, defaults to ``False``.
    r   formatzTuple[int, int]spatial_sizeztorch.Tensor)tensorr   r   returnc                C  s   | | }||_||_|S )N)as_subclassr   r   )clsr   r   r   Zbounding_boxr   r   r   _wrap-   s    
zBoundingBox._wrapNdtypedevicerequires_gradr   zUnion[BoundingBoxFormat, str]zOptional[torch.dtype]z'Optional[Union[torch.device, str, int]]zOptional[bool])datar   r   r!   r"   r#   r   c                C  s8   | j ||||d}t|tr(t|  }| j|||dS )Nr    r   r   )Z
_to_tensor
isinstancestrr   upperr   )r   r$   r   r   r!   r"   r#   r   r   r   r   __new__4   s    

zBoundingBox.__new__r%   zOptional[BoundingBoxFormat]zOptional[Tuple[int, int]])otherr   r   r   r   c                C  sB   t |trt|  }| j||dur(|n|j|dur8|n|jdS )a  Wrap a :class:`torch.Tensor` as :class:`BoundingBox` from a reference.

        Args:
            other (BoundingBox): Reference bounding box.
            tensor (Tensor): Tensor to be wrapped as :class:`BoundingBox`
            format (BoundingBoxFormat, str, optional): Format of the bounding box.  If omitted, it is taken from the
                reference.
            spatial_size (two-tuple of ints, optional): Height and width of the corresponding image or video. If
                omitted, it is taken from the reference.

        Nr%   )r&   r'   r   r(   r   r   r   )r   r*   r   r   r   r   r   r   	wrap_likeE   s    
zBoundingBox.wrap_like)tensor_contentsr'   )r,   r   c                C  s   | j | j| jdS Nr%   )Z
_make_reprr   r   )selfr,   r   r   r   __repr__b   s    zBoundingBox.__repr__)r   c                 C  s*   | j j| tj| j| jd}t| |S r-   )	_FZhorizontal_flip_bounding_boxr   torchTensorr   r   r   r+   r.   outputr   r   r   horizontal_flipe   s    zBoundingBox.horizontal_flipc                 C  s*   | j j| tj| j| jd}t| |S r-   )	r0   Zvertical_flip_bounding_boxr   r1   r2   r   r   r   r+   r3   r   r   r   vertical_flipk   s    zBoundingBox.vertical_flipwarnz	List[int]zUnion[InterpolationMode, int]zOptional[int]zOptional[Union[str, bool]])sizeinterpolationmax_size	antialiasr   c                 C  s2   | j j| tj| j||d\}}tj| ||dS )N)r   r8   r:   r   )r0   Zresize_bounding_boxr   r1   r2   r   r   r+   )r.   r8   r9   r:   r;   r4   r   r   r   r   resizeq   s    

zBoundingBox.resizeint)topleftheightwidthr   c                 C  s6   | j j| tj| j||||d\}}tj| ||dS )N)r?   r@   rA   rB   r<   )r0   Zcrop_bounding_boxr   r1   r2   r   r   r+   )r.   r?   r@   rA   rB   r4   r   r   r   r   crop   s    
zBoundingBox.crop)output_sizer   c                 C  s4   | j j| tj| j| j|d\}}tj| ||dS )N)r   r   rD   r<   )	r0   Zcenter_crop_bounding_boxr   r1   r2   r   r   r   r+   )r.   rD   r4   r   r   r   r   center_crop   s    
zBoundingBox.center_crop)r?   r@   rA   rB   r8   r9   r;   r   c           
   	   C  s8   | j j| tj| j|||||d\}}	tj| ||	dS )N)r8   r<   )r0   Zresized_crop_bounding_boxr   r1   r2   r   r   r+   )
r.   r?   r@   rA   rB   r8   r9   r;   r4   r   r   r   r   resized_crop   s    

zBoundingBox.resized_cropconstantzUnion[int, Sequence[int]]z(Optional[Union[int, float, List[float]]])paddingfillpadding_moder   c                 C  s6   | j j| tj| j| j||d\}}tj| ||dS )N)r   r   rH   rJ   r<   )	r0   Zpad_bounding_boxr   r1   r2   r   r   r   r+   )r.   rH   rI   rJ   r4   r   r   r   r   pad   s    

zBoundingBox.padFfloatboolzOptional[List[float]]r   )angler9   expandcenterrI   r   c                 C  s8   | j j| tj| j| j|||d\}}tj| ||dS )N)r   r   rN   rO   rP   r<   )	r0   Zrotate_bounding_boxr   r1   r2   r   r   r   r+   )r.   rN   r9   rO   rP   rI   r4   r   r   r   r   rotate   s    

zBoundingBox.rotatezUnion[int, float]zList[float])rN   	translatescaleshearr9   rI   rP   r   c           	   
   C  s4   | j j| tj| j| j|||||d}t| |S )N)rR   rS   rT   rP   )	r0   Zaffine_bounding_boxr   r1   r2   r   r   r   r+   )	r.   rN   rR   rS   rT   r9   rI   rP   r4   r   r   r   affine   s    


zBoundingBox.affinezOptional[List[List[int]]])startpoints	endpointsr9   rI   coefficientsr   c                 C  s0   | j j| tj| j| j|||d}t| |S )N)r   r   rV   rW   rX   )	r0   Zperspective_bounding_boxr   r1   r2   r   r   r   r+   )r.   rV   rW   r9   rI   rX   r4   r   r   r   perspective   s    
zBoundingBox.perspective)displacementr9   rI   r   c                 C  s,   | j j| tj| j| j|d}t| |S )N)rZ   )	r0   Zelastic_bounding_boxr   r1   r2   r   r   r   r+   )r.   rZ   r9   rI   r4   r   r   r   elastic   s    zBoundingBox.elastic)NrG   )r   r   r   r   __annotations__classmethodr   r)   r+   r/   r5   r6   r
   ZBILINEARr=   rC   rE   rF   rK   ZNEARESTrQ   rU   rY   r[   r   r   r   r   r      sX   
	  r   )
__future__r   enumr   typingr   r   r   r   r   r	   r1   Ztorchvision.transformsr
   Z
_datapointr   r   r   r   r   r   r   r   <module>   s    