U
    hZ                     @  s@  d dl mZ d dlZd dlmZmZmZmZmZm	Z	 d dl
Z
d dlZd dlmZmZ d dlmZmZmZ d dlmZ d dlm  m  mZ d dlmZmZmZ d dlm Z  d d	l!m"Z" d d
l#m$Z$m%Z%m&Z& d dl'm(Z(m)Z)m*Z* d dl+m,Z, ddddgZ-dZ.G dd de*Z/G dd de*Z0G dd de*Z1G dd de*Z2dS )    )annotationsN)AnyCallableLiteralSequenceTuplecast)add_weightedget_num_channels)AfterValidatorFieldfield_validator)	Annotated)adapt_pixel_distributionapply_histogramfourier_domain_adaptation)read_rgb_image)Compose)ZeroOneRangeTypecheck_01nondecreasing)BaseTransformInitSchemaBasicTransformImageOnlyTransform)ScaleFloatTypeHistogramMatchingFDAPixelDistributionAdaptationTemplateTransform      ?c                      s   e Zd ZdZG dd deZdeddfddd	d
dd fddZdddddddddZddddZ	ddddZ
ddddZ  ZS )r   a
  Adjust the pixel values of an input image to match the histogram of a reference image.

    This transform applies histogram matching, a technique that modifies the distribution of pixel
    intensities in the input image to closely resemble that of a reference image. This process is
    performed independently for each channel in multi-channel images, provided both the input and
    reference images have the same number of channels.

    Histogram matching is particularly useful for:
    - Normalizing images from different sources or captured under varying conditions.
    - Preparing images for feature matching or other computer vision tasks where consistent
      tone and contrast are important.
    - Simulating different lighting or camera conditions in a controlled manner.

    Args:
        reference_images (Sequence[Any]): A sequence of reference image sources. These can be
            file paths, URLs, or any objects that can be converted to images by the `read_fn`.
        blend_ratio (tuple[float, float]): Range for the blending factor between the original
            and the matched image. Must be two floats between 0 and 1, where:
            - 0 means no blending (original image is returned)
            - 1 means full histogram matching
            A random value within this range is chosen for each application.
            Default: (0.5, 1.0)
        read_fn (Callable[[Any], np.ndarray]): A function that takes an element from
            `reference_images` and returns a numpy array representing the image.
            Default: read_rgb_image (reads image file from disk)
        p (float): Probability of applying the transform. Default: 0.5

    Targets:
        image

    Image types:
        uint8, float32

    Note:
        - This transform cannot be directly serialized due to its dependency on external image data.
        - The effectiveness of the matching depends on the similarity between the input and reference images.
        - For best results, choose reference images that represent the desired tone and contrast.

    Example:
        >>> import numpy as np
        >>> import albumentations as A
        >>> image = np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)
        >>> reference_image = np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)
        >>> transform = A.HistogramMatching(
        ...     reference_images=[reference_image],
        ...     blend_ratio=(0.5, 1.0),
        ...     read_fn=lambda x: x,
        ...     p=1
        ... )
        >>> result = transform(image=image)
        >>> matched_image = result["image"]

    References:
        - Histogram Matching in scikit-image:
          https://scikit-image.org/docs/dev/auto_examples/color_exposure/plot_histogram_matching.html
    c                   @  s&   e Zd ZU ded< ded< ded< dS )zHistogramMatching.InitSchemaSequence[Any]reference_imagesWAnnotated[tuple[float, float], AfterValidator(nondecreasing), AfterValidator(check_01)]blend_ratioCallable[[Any], np.ndarray]read_fnN__name__
__module____qualname____annotations__ r+   r+   ]/tmp/pip-unpacked-wheel-e8onvpoz/albumentations/augmentations/domain_adaptation/transforms.py
InitSchema\   s   
r-   )r         ?Nr   r    tuple[float, float]r$   bool | Nonefloat)r!   r#   r%   always_applypc                   s&   t  j||d || _|| _|| _d S N)r3   r2   )super__init__r!   r%   r#   )selfr!   r#   r%   r2   r3   	__class__r+   r,   r6   a   s    zHistogramMatching.__init__
np.ndarrayr   )r7   imgreference_imager#   paramsreturnc                 K  s   t |||S N)r   r7   r;   r<   r#   r=   r+   r+   r,   applyn   s    zHistogramMatching.applydict[str, np.ndarray]r>   c                 C  s    |  t| jtj| j dS N)r<   r#   r%   randomchoicer!   uniformr#   r7   r+   r+   r,   
get_paramsw   s    
zHistogramMatching.get_paramsztuple[str, ...]c                 C  s   dS )N)r!   r#   r%   r+   rI   r+   r+   r,   get_transform_init_args_names}   s    z/HistogramMatching.get_transform_init_args_namesdict[str, Any]c                 C  s   d}t |d S )Nz(HistogramMatching can not be serialized.NotImplementedErrorr7   msgr+   r+   r,   to_dict_private   s    z!HistogramMatching.to_dict_privater'   r(   r)   __doc__r   r-   r   r6   rA   rJ   rK   rQ   __classcell__r+   r+   r8   r,   r   "   s   9	c                      s   e Zd ZdZG dd deZdeddfddd	d
dd fddZddddddddZddddddZ	ddddZ
ddddZddd d!Z  ZS )"r   a  Fourier Domain Adaptation (FDA) for simple "style transfer" in the context of unsupervised domain adaptation
    (UDA). FDA manipulates the frequency components of images to reduce the domain gap between source
    and target datasets, effectively adapting images from one domain to closely resemble those from another without
    altering their semantic content.

    This transform is particularly beneficial in scenarios where the training (source) and testing (target) images
    come from different distributions, such as synthetic versus real images, or day versus night scenes.
    Unlike traditional domain adaptation methods that may require complex adversarial training, FDA achieves domain
    alignment by swapping low-frequency components of the Fourier transform between the source and target images.
    This technique has shown to improve the performance of models on the target domain, particularly for tasks
    like semantic segmentation, without additional training for domain invariance.

    The 'beta_limit' parameter controls the extent of frequency component swapping, with lower values preserving more
    of the original image's characteristics and higher values leading to more pronounced adaptation effects.
    It is recommended to use beta values less than 0.3 to avoid introducing artifacts.

    Args:
        reference_images (Sequence[Any]): Sequence of objects to be converted into images by `read_fn`. This typically
            involves paths to images that serve as target domain examples for adaptation.
        beta_limit (tuple[float, float] | float): Coefficient beta from the paper, controlling the swapping extent of
            frequency components. If one value is provided beta will be sampled from uniform
            distribution [0, beta_limit]. Values should be less than 0.5.
        read_fn (Callable): User-defined function for reading images. It takes an element from `reference_images` and
            returns a numpy array of image pixels. By default, it is expected to take a path to an image and return a
            numpy array.

    Targets:
        image

    Image types:
        uint8, float32

    Reference:
        - https://github.com/YanchaoYang/FDA
        - https://openaccess.thecvf.com/content_CVPR_2020/papers/Yang_FDA_Fourier_Domain_Adaptation_for_Semantic_Segmentation_CVPR_2020_paper.pdf

    Example:
        >>> import numpy as np
        >>> import albumentations as A
        >>> image = np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)
        >>> target_image = np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)
        >>> aug = A.Compose([A.FDA([target_image], p=1, read_fn=lambda x: x)])
        >>> result = aug(image=image)

    Note:
        FDA is a powerful tool for domain adaptation, particularly in unsupervised settings where annotated target
        domain samples are unavailable. It enables significant improvements in model generalization by aligning
        the low-level statistics of source and target images through a simple yet effective Fourier-based method.
    c                   @  sB   e Zd ZU ded< ded< ded< ededddd	d
ZdS )zFDA.InitSchemar    r!   r$   r%   r   
beta_limitr/   )valuer>   c                 C  sR   dt f}|d |d   kr6|d   kr6|d ksNn td| d| d|S )Nr      zValues should be in the range z got  )MAX_BETA_LIMIT
ValueError)clsrV   Zboundsr+   r+   r,   check_ranges   s    0zFDA.InitSchema.check_rangesN)r'   r(   r)   r*   r   classmethodr\   r+   r+   r+   r,   r-      s   
r-   )r   g?Nr   r    r   r$   r0   r1   )r!   rU   r%   r2   r3   c                   s4   t  j||d || _|| _ttttf || _d S r4   )r5   r6   r!   r%   r   r   r1   rU   )r7   r!   rU   r%   r2   r3   r8   r+   r,   r6      s    zFDA.__init__r:   r   )r;   target_imagebetar=   r>   c                 K  s   t |||S r?   )r   )r7   r;   r^   r_   r=   r+   r+   r,   rA      s    z	FDA.applyrL   rB   r=   datar>   c                 C  s4   |  t| j}tj||d |d fd}d|iS )NcolsZrows)Zdsizer^   )r%   rF   rG   r!   cv2resize)r7   r=   ra   Z
target_imgr+   r+   r,   get_params_dependent_on_data   s    z FDA.get_params_dependent_on_datadict[str, float]rC   c                 C  s   dt j| j iS )Nr_   )rF   rH   rU   rI   r+   r+   r,   rJ      s    zFDA.get_paramsztuple[str, str, str]c                 C  s   dS )N)r!   rU   r%   r+   rI   r+   r+   r,   rK      s    z!FDA.get_transform_init_args_namesc                 C  s   d}t |d S )NzFDA can not be serialized.rM   rO   r+   r+   r,   rQ      s    zFDA.to_dict_private)r'   r(   r)   rS   r   r-   r   r6   rA   re   rJ   rK   rQ   rT   r+   r+   r8   r,   r      s   2	c                      s   e Zd ZdZG dd deZdedddfdd	d
dddd fddZddddddddZddddZ	ddddZ
ddddZ  ZS )r   ac  Performs pixel-level domain adaptation by aligning the pixel value distribution of an input image
    with that of a reference image. This process involves fitting a simple statistical transformation
    (such as PCA, StandardScaler, or MinMaxScaler) to both the original and the reference images,
    transforming the original image with the transformation trained on it, and then applying the inverse
    transformation using the transform fitted on the reference image. The result is an adapted image
    that retains the original content while mimicking the pixel value distribution of the reference domain.

    The process can be visualized as two main steps:
    1. Adjusting the original image to a standard distribution space using a selected transform.
    2. Moving the adjusted image into the distribution space of the reference image by applying the inverse
       of the transform fitted on the reference image.

    This technique is especially useful in scenarios where images from different domains (e.g., synthetic
    vs. real images, day vs. night scenes) need to be harmonized for better consistency or performance in
    image processing tasks.

    Args:
        reference_images (Sequence[Any]): A sequence of objects (typically image paths) that will be
            converted into images by `read_fn`. These images serve as references for the domain adaptation.
        blend_ratio (tuple[float, float]): Specifies the minimum and maximum blend ratio for mixing
            the adapted image with the original. This enhances the diversity of the output images.
            Values should be in the range [0, 1]. Default: (0.25, 1.0)
        read_fn (Callable): A user-defined function for reading and converting the objects in
            `reference_images` into numpy arrays. By default, it assumes these objects are image paths.
        transform_type (Literal["pca", "standard", "minmax"]): Specifies the type of statistical
            transformation to apply.
            - "pca": Principal Component Analysis
            - "standard": StandardScaler (zero mean and unit variance)
            - "minmax": MinMaxScaler (scales to a fixed range, usually [0, 1])
            Default: "pca"
        p (float): The probability of applying the transform to any given image. Default: 0.5

    Targets:
        image

    Image types:
        uint8, float32

    Number of channels:
        Any

    Note:
        - The effectiveness of the adaptation depends on the similarity between the input and reference domains.
        - PCA transformation may alter color relationships more significantly than other methods.
        - StandardScaler and MinMaxScaler preserve color relationships better but may provide less dramatic adaptations.
        - The blend_ratio parameter allows for a smooth transition between the original and fully adapted image.
        - This transform cannot be directly serialized due to its dependency on external image data.

    Example:
        >>> import numpy as np
        >>> import albumentations as A
        >>> image = np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)
        >>> reference_image = np.random.randint(0, 256, [100, 100, 3], dtype=np.uint8)
        >>> transform = A.PixelDistributionAdaptation(
        ...     reference_images=[reference_image],
        ...     blend_ratio=(0.5, 1.0),
        ...     transform_type="standard",
        ...     read_fn=lambda x: x,
        ...     p=1.0
        ... )
        >>> result = transform(image=image)
        >>> adapted_image = result["image"]

    References:
        - https://github.com/arsenyinfo/qudida
        - https://arxiv.org/abs/1911.11483
    c                   @  s.   e Zd ZU ded< ded< ded< ded< d	S )
z&PixelDistributionAdaptation.InitSchemar    r!   r"   r#   r$   r%   &Literal[('pca', 'standard', 'minmax')]transform_typeNr&   r+   r+   r+   r,   r-   1  s   
r-   )g      ?r.   ZpcaNr   r    r/   r$   rg   r0   r1   )r!   r#   r%   rh   r2   r3   c                   s,   t  j||d || _|| _|| _|| _d S r4   )r5   r6   r!   r%   r#   rh   )r7   r!   r#   r%   rh   r2   r3   r8   r+   r,   r6   7  s
    	z$PixelDistributionAdaptation.__init__r:   r   )r;   r<   r#   r=   r>   c                 K  s   t |||| jdS )N)refZweightrh   )r   rh   r@   r+   r+   r,   rA   F  s    z!PixelDistributionAdaptation.applyrL   rC   c                 C  s    |  t| jtj| j dS rD   rE   rI   r+   r+   r,   rJ   N  s    
z&PixelDistributionAdaptation.get_paramsztuple[str, str, str, str]c                 C  s   dS )N)r!   r#   r%   rh   r+   rI   r+   r+   r,   rK   T  s    z9PixelDistributionAdaptation.get_transform_init_args_namesc                 C  s   d}t |d S )Nz2PixelDistributionAdaptation can not be serialized.rM   rO   r+   r+   r,   rQ   W  s    z+PixelDistributionAdaptation.to_dict_privaterR   r+   r+   r8   r,   r      s   D	c                	      s   e Zd ZdZG dd deZd#ddd	d
dddd fddZddddddddZddddZddddddZ	e
dddd Zddd!d"Z  ZS )$r   ap  Apply blending of input image with specified templates.

    This transform overlays one or more template images onto the input image using alpha blending.
    It allows for creating complex composite images or simulating various visual effects.

    Args:
        templates (numpy array | list[np.ndarray]): Images to use as templates for the transform.
            If a single numpy array is provided, it will be used as the only template.
            If a list of numpy arrays is provided, one will be randomly chosen for each application.

        img_weight (tuple[float, float]  | float): Weight of the original image in the blend.
            If a single float, that value will always be used.
            If a tuple (min, max), the weight will be randomly sampled from the range [min, max) for each application.
            To use a fixed weight, use (weight, weight).
            Default: (0.5, 0.5).

        template_transform (A.Compose | None): A composition of Albumentations transforms to apply to the template
            before blending.
            This should be an instance of A.Compose containing one or more Albumentations transforms.
            Default: None.

        name (str | None): Name of the transform instance. Used for serialization purposes.
            Default: None.

        p (float): Probability of applying the transform. Default: 0.5.

    Targets:
        image

    Image types:
        uint8, float32

    Number of channels:
        Any

    Note:
        - The template(s) must have the same number of channels as the input image or be single-channel.
        - If a single-channel template is used with a multi-channel image, the template will be replicated across
          all channels.
        - The template(s) will be resized to match the input image size if they differ.
        - To make this transform serializable, provide a name when initializing it.

    Mathematical Formulation:
        Given:
        - I: Input image
        - T: Template image
        - w_i: Weight of input image (sampled from img_weight)

        The blended image B is computed as:

        B = w_i * I + (1 - w_i) * T

    Examples:
        >>> import numpy as np
        >>> import albumentations as A
        >>> image = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)
        >>> template = np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8)

        # Apply template transform with a single template
        >>> transform = A.TemplateTransform(templates=template, name="my_template_transform", p=1.0)
        >>> blended_image = transform(image=image)['image']

        # Apply template transform with multiple templates and custom weights
        >>> templates = [np.random.randint(0, 256, (100, 100, 3), dtype=np.uint8) for _ in range(3)]
        >>> transform = A.TemplateTransform(
        ...     templates=templates,
        ...     img_weight=(0.3, 0.7),
        ...     name="multi_template_transform",
        ...     p=1.0
        ... )
        >>> blended_image = transform(image=image)['image']

        # Apply template transform with additional transforms on the template
        >>> template_transform = A.Compose([A.RandomBrightnessContrast(p=1)])
        >>> transform = A.TemplateTransform(
        ...     templates=template,
        ...     img_weight=0.6,
        ...     template_transform=template_transform,
        ...     name="transformed_template",
        ...     p=1.0
        ... )
        >>> blended_image = transform(image=image)['image']

    References:
        - Alpha compositing: https://en.wikipedia.org/wiki/Alpha_compositing
        - Image blending: https://en.wikipedia.org/wiki/Image_blending
    c                   @  s`   e Zd ZU ded< ded< eddZded< d	Zd
ed< ded< ededddddZ	d	S )zTemplateTransform.InitSchemaz!np.ndarray | Sequence[np.ndarray]	templatesr   
img_weightzITemplate_weight is deprecated. Computed automatically as (1 - img_weight))
deprecatedzZeroOneRangeType | Nonetemplate_weightNCompose | BasicTransform | Nonetemplate_transform
str | Nonenamenp.ndarray | list[np.ndarray]zlist[np.ndarray])vr>   c                 C  sN   t |tjr|gS t |tr>tdd |D s:d}t||S d}t|d S )Nc                 s  s   | ]}t |tjV  qd S r?   )
isinstancenpndarray).0itemr+   r+   r,   	<genexpr>  s     zBTemplateTransform.InitSchema.validate_templates.<locals>.<genexpr>z#All templates must be numpy arrays.z:Templates must be a numpy array or a list of numpy arrays.)rt   ru   rv   listallrZ   	TypeError)r[   rs   rP   r+   r+   r,   validate_templates  s    
z/TemplateTransform.InitSchema.validate_templates)
r'   r(   r)   r*   r   rm   ro   r   r]   r}   r+   r+   r+   r,   r-     s   
r-   r   r   Nr   rr   r   Nonern   rp   r0   r1   )rj   rk   rm   ro   rq   r2   r3   c                   s:   t  j||d || _ttttf || _|| _|| _d S r4   )	r5   r6   rj   r   r   r1   rk   ro   rq   )r7   rj   rk   rm   ro   rq   r2   r3   r8   r+   r,   r6     s
    
zTemplateTransform.__init__r:   r   )r;   templaterk   r=   r>   c                 K  s*   |dkr|S |dkr|S t |||d| S )Nr   rW   )r	   )r7   r;   r   rk   r=   r+   r+   r,   rA     s
    zTemplateTransform.applyrf   rC   c                 C  s   dt j| j iS )Nrk   )rF   rH   rk   rI   r+   r+   r,   rJ     s     
zTemplateTransform.get_paramsrL   r`   c                 C  s
  d|kr|d n
|d d }t | j}| jd k	rB| j|dd }t|dt|fkrvdt| dt| }t||j|jkrd}t||jd d	 |jd d	 krtj	||jd d	 t
jd
}t|dkrt|dkrtj|ft| dd}||j}d|iS )NimageZimagesr   )r   rW   zUTemplate must be a single channel or has the same number of channels as input image (z), got z.Image and template must be the same image type   )interpolation)Zaxisr   )rF   rG   rj   ro   r
   rZ   Zdtypeshape
fgeometricrd   rc   Z
INTER_AREAru   stackZreshape)r7   r=   ra   r;   r   rP   r+   r+   r,   re     s"    
z.TemplateTransform.get_params_dependent_on_databoolc                 C  s   dS )NFr+   )r[   r+   r+   r,   is_serializable  s    z!TemplateTransform.is_serializablec                 C  s&   | j d krd}t||  | j dS )NzTo make a TemplateTransform serializable you should provide the `name` argument, e.g. `TemplateTransform(name='my_transform', ...)`.)Z__class_fullname__r'   )rq   rZ   Zget_class_fullnamerO   r+   r+   r,   rQ     s
    
z!TemplateTransform.to_dict_private)r~   NNNNr   )r'   r(   r)   rS   r   r-   r6   rA   rJ   re   r]   r   rQ   rT   r+   r+   r8   r,   r   \  s   X       )3
__future__r   rF   typingr   r   r   r   r   r   rc   Znumpyru   Zalbucorer	   r
   Zpydanticr   r   r   Ztyping_extensionsr   Z1albumentations.augmentations.geometric.functionalZaugmentationsZ	geometricZ
functionalr   Z9albumentations.augmentations.domain_adaptation.functionalr   r   r   Z"albumentations.augmentations.utilsr   Zalbumentations.core.compositionr   Zalbumentations.core.pydanticr   r   r   Z(albumentations.core.transforms_interfacer   r   r   Zalbumentations.core.typesr   __all__rY   r   r   r   r   r+   r+   r+   r,   <module>   s0    cgp