o oiZ@s|ddlmZmZddlZddlmmZgdZdej dej fddZ    d!d ej dej d eej deee de de de de dej fddZ    d!d ej dej d eej deee de de de de dej fddZ    d!d ej dej d eej deee de de de de dej fddZ    d!d ej dej d eej deee de de de de dej fddZ    d!d ej dej d eej deee de de de de dej fddZ    d!d ej dej d eej deee de de de de dej fddZ    d!d ej dej d eej deee de de de de dej fdd ZdS)")ListOptionalN) bottom_hatclosingdilationerosiongradientopeningtop_hatkernelreturncCs8|\}}tj|||j|jd}|||d||S)N)dtypedevice)sizetorcheyer rview)r hwrP/home/ubuntu/.local/lib/python3.10/site-packages/kornia/morphology/morphology.py_neight2channels_like_kernels rgeodesic@unfoldtensorstructuring_elementorigin border_type border_valuemax_valenginecCs t|tjstdt|t|jdkrtd|t|tjs.tdt|t|jdkr>td||j\}} |durO|d| dg}|d| |dd|d ||d dg} |d krn| }d }t j || ||d } |durt |} | | |d k<n | } | | |d k<|d kr| d|d d| d} t| | dd\} } t| d\} } nG|dkr|\}}}}| jdd\}}t|}t j| ||d|||d | dd djdd\} } | ||||} ntd|d| |S)aReturn the dilated image applying the same kernel in each channel. .. image:: _static/img/dilation.png The kernel must have 2 dimensions. Args: tensor: Image with shape :math:`(B, C, H, W)`. kernel: Positions of non-infinite elements of a flat structuring element. Non-zero values give the set of neighbors of the center over which the operation is applied. Its shape is :math:`(k_x, k_y)`. For full structural elements use torch.ones_like(structural_element). structuring_element: Structuring element used for the grayscale dilation. It may be a non-flat structuring element. origin: Origin of the structuring element. Default: ``None`` and uses the center of the structuring element as origin (rounding towards zero). border_type: It determines how the image borders are handled, where ``border_value`` is the value when ``border_type`` is equal to ``constant``. Default: ``geodesic`` which ignores the values that are outside the image when applying the operation. border_value: Value to fill past edges of input if ``border_type`` is ``constant``. max_val: The value of the infinite elements in the kernel. engine: convolution is faster and less memory hungry, and unfold is more stable numerically Returns: Dilated image with shape :math:`(B, C, H, W)`. .. note:: See a working example `here `__. Example: >>> tensor = torch.rand(1, 3, 5, 5) >>> kernel = torch.ones(3, 3) >>> dilated_img = dilation(tensor, kernel) &Input type is not a torch.Tensor. Got 'Input size must have 4 dimensions. Got 'Kernel type is not a torch.Tensor. Got (Kernel size must have 2 dimensions. Got Nrrrconstantmodevaluer)rr convolutionpaddingbiasdimengine * is unknown, use 'convolution' or 'unfold') isinstancerTensor TypeErrortypelenshape ValueErrorr6Fpad zeros_likeclonermaxfliprrconv2drNotImplementedErrorview_as)rr rrr r!r"r#se_hse_wpad_eoutput neighborhood_BCHWh_padw_padreshape_kernelrrrr sJ ,  , $  rcCst|tjstdt|t|jdkrtd|t|tjs.tdt|t|jdkr>td||j\}} |durO|d| dg}|d| |dd|d ||d dg} |d krm|}d }t j || ||d } |durt |} | | |d k<n | } | | |d k<|d kr| d|d d| d} t| | d\} } t| d\} } | S|dkr|\}}}}| jdd\}}t|}t j| ||d|||d | d djdd\} } | ||||} | Std|d)aReturn the eroded image applying the same kernel in each channel. .. image:: _static/img/erosion.png The kernel must have 2 dimensions. Args: tensor: Image with shape :math:`(B, C, H, W)`. kernel: Positions of non-infinite elements of a flat structuring element. Non-zero values give the set of neighbors of the center over which the operation is applied. Its shape is :math:`(k_x, k_y)`. For full structural elements use torch.ones_like(structural_element). structuring_element (torch.Tensor, optional): Structuring element used for the grayscale dilation. It may be a non-flat structuring element. origin: Origin of the structuring element. Default: ``None`` and uses the center of the structuring element as origin (rounding towards zero). border_type: It determines how the image borders are handled, where ``border_value`` is the value when ``border_type`` is equal to ``constant``. Default: ``geodesic`` which ignores the values that are outside the image when applying the operation. border_value: Value to fill past edges of input if border_type is ``constant``. max_val: The value of the infinite elements in the kernel. engine: ``convolution`` is faster and less memory hungry, and ``unfold`` is more stable numerically Returns: Eroded image with shape :math:`(B, C, H, W)`. .. note:: See a working example `here `__. Example: >>> tensor = torch.rand(1, 3, 5, 5) >>> kernel = torch.ones(5, 5) >>> output = erosion(tensor, kernel) r$r%r&r'r(r)Nrrrr*r+rr.r/r0r1r2r5r7r8)r9rr:r;r<r=r>r?r6r@rArBrCrminrrrFrrG)rr rrr r!r"r#rIrJrKrLrMrNrOrPrQrRHpadWpadrUrrrr}sL ,  ,    rc Ct|tjstdt|t|jdkrtd|t|tjs.tdt|t|jdkr>td|t t ||||||||d|||||||dS)aReturn the opened image, (that means, dilation after an erosion) applying the same kernel in each channel. .. image:: _static/img/opening.png The kernel must have 2 dimensions. Args: tensor: Image with shape :math:`(B, C, H, W)`. kernel: Positions of non-infinite elements of a flat structuring element. Non-zero values give the set of neighbors of the center over which the operation is applied. Its shape is :math:`(k_x, k_y)`. For full structural elements use torch.ones_like(structural_element). structuring_element: Structuring element used for the grayscale dilation. It may be a non-flat structuring element. origin: Origin of the structuring element. Default: ``None`` and uses the center of the structuring element as origin (rounding towards zero). border_type: It determines how the image borders are handled, where ``border_value`` is the value when ``border_type`` is equal to ``constant``. Default: ``geodesic`` which ignores the values that are outside the image when applying the operation. border_value: Value to fill past edges of input if ``border_type`` is ``constant``. max_val: The value of the infinite elements in the kernel. engine: convolution is faster and less memory hungry, and unfold is more stable numerically Returns: torch.Tensor: Opened image with shape :math:`(B, C, H, W)`. .. note:: See a working example `here `__. Example: >>> tensor = torch.rand(1, 3, 5, 5) >>> kernel = torch.ones(3, 3) >>> opened_img = opening(tensor, kernel) r$r%r&r'r(r)r rrr r!r"r#) r9rr:r;r<r=r>r?r6rrrr rrr r!r"r#rrrr 6 ,  r c CrY)aReturn the closed image, (that means, erosion after a dilation) applying the same kernel in each channel. .. image:: _static/img/closing.png The kernel must have 2 dimensions. Args: tensor: Image with shape :math:`(B, C, H, W)`. kernel: Positions of non-infinite elements of a flat structuring element. Non-zero values give the set of neighbors of the center over which the operation is applied. Its shape is :math:`(k_x, k_y)`. For full structural elements use torch.ones_like(structural_element). structuring_element: Structuring element used for the grayscale dilation. It may be a non-flat structuring element. origin: Origin of the structuring element. Default is None and uses the center of the structuring element as origin (rounding towards zero). border_type: It determines how the image borders are handled, where ``border_value`` is the value when ``border_type`` is equal to ``constant``. Default: ``geodesic`` which ignores the values that are outside the image when applying the operation. border_value: Value to fill past edges of input if ``border_type`` is ``constant``. max_val: The value of the infinite elements in the kernel. engine: convolution is faster and less memory hungry, and unfold is more stable numerically Returns: Closed image with shape :math:`(B, C, H, W)`. .. note:: See a working example `here `__. Example: >>> tensor = torch.rand(1, 3, 5, 5) >>> kernel = torch.ones(3, 3) >>> closed_img = closing(tensor, kernel) r$r%r&r'r(r)rZ) r9rr:r;r<r=r>r?r6rrr[rrrr(r\rc Cs0t||||||||dt||||||||dS)aReturn the morphological gradient of an image. .. image:: _static/img/gradient.png That means, (dilation - erosion) applying the same kernel in each channel. The kernel must have 2 dimensions. Args: tensor: Image with shape :math:`(B, C, H, W)`. kernel: Positions of non-infinite elements of a flat structuring element. Non-zero values give the set of neighbors of the center over which the operation is applied. Its shape is :math:`(k_x, k_y)`. For full structural elements use torch.ones_like(structural_element). structuring_element: Structuring element used for the grayscale dilation. It may be a non-flat structuring element. origin: Origin of the structuring element. Default is None and uses the center of the structuring element as origin (rounding towards zero). border_type: It determines how the image borders are handled, where ``border_value`` is the value when ``border_type`` is equal to ``constant``. Default: ``geodesic`` which ignores the values that are outside the image when applying the operation. border_value: Value to fill past edges of input if ``border_type`` is ``constant``. max_val: The value of the infinite elements in the kernel. engine: convolution is faster and less memory hungry, and unfold is more stable numerically Returns: Gradient image with shape :math:`(B, C, H, W)`. .. note:: See a working example `here `__. Example: >>> tensor = torch.rand(1, 3, 5, 5) >>> kernel = torch.ones(3, 3) >>> gradient_img = gradient(tensor, kernel) rZ)rrr[rrrrvs*- rc Cst|tjstdt|t|jdkrtd|t|tjs.tdt|t|jdkr>td||t ||||||||dS)aReturn the top hat transformation of an image. .. image:: _static/img/top_hat.png That means, (image - opened_image) applying the same kernel in each channel. The kernel must have 2 dimensions. See :func:`~kornia.morphology.opening` for details. Args: tensor: Image with shape :math:`(B, C, H, W)`. kernel: Positions of non-infinite elements of a flat structuring element. Non-zero values give the set of neighbors of the center over which the operation is applied. Its shape is :math:`(k_x, k_y)`. For full structural elements use torch.ones_like(structural_element). structuring_element: Structuring element used for the grayscale dilation. It may be a non-flat structuring element. origin: Origin of the structuring element. Default: ``None`` and uses the center of the structuring element as origin (rounding towards zero). border_type: It determines how the image borders are handled, where ``border_value`` is the value when ``border_type`` is equal to ``constant``. Default: ``geodesic`` which ignores the values that are outside the image when applying the operation. border_value: Value to fill past edges of input if ``border_type`` is ``constant``. max_val: The value of the infinite elements in the kernel. engine: convolution is faster and less memory hungry, and unfold is more stable numerically Returns: Top hat transformed image with shape :math:`(B, C, H, W)`. .. note:: See a working example `here `__. Example: >>> tensor = torch.rand(1, 3, 5, 5) >>> kernel = torch.ones(3, 3) >>> top_hat_img = top_hat(tensor, kernel) r$r%r&r'r(r)rZ) r9rr:r;r<r=r>r?r6r r[rrrr s$ / r c Cst|tjstdt|t|jdkrtd|t|tjs.tdt|t|jdkr>td|t ||||||||d|S)a%Return the bottom hat transformation of an image. .. image:: _static/img/bottom_hat.png That means, (closed_image - image) applying the same kernel in each channel. The kernel must have 2 dimensions. See :func:`~kornia.morphology.closing` for details. Args: tensor: Image with shape :math:`(B, C, H, W)`. kernel: Positions of non-infinite elements of a flat structuring element. Non-zero values give the set of neighbors of the center over which the operation is applied. Its shape is :math:`(k_x, k_y)`. For full structural elements use torch.ones_like(structural_element). structuring_element: Structuring element used for the grayscale dilation. It may be a non-flat structuring element. origin: Origin of the structuring element. Default: ``None`` and uses the center of the structuring element as origin (rounding towards zero). border_type: It determines how the image borders are handled, where ``border_value`` is the value when ``border_type`` is equal to ``constant``. Default: ``geodesic`` which ignores the values that are outside the image when applying the operation. border_value: Value to fill past edges of input if ``border_type`` is ``constant``. max_val: The value of the infinite elements in the kernel. engine: convolution is faster and less memory hungry, and unfold is more stable numerically Returns: Top hat transformed image with shape :math:`(B, C, H, W)`. .. note:: See a working example `here `__. Example: >>> tensor = torch.rand(1, 3, 5, 5) >>> kernel = torch.ones(3, 3) >>> bottom_hat_img = bottom_hat(tensor, kernel) r$r%r&r'r(r)rZ) r9rr:r;r<r=r>r?r6rr[rrrrs* /  r)NNrrrr)typingrrrtorch.nn.functionalnn functionalr@__all__r:rintstrfloatrrr rrr rrrrrsh   `  a  P  Q  E  J