o
    i&                     @  s   d dl mZ d dlmZmZ d dlZd dlZd dlm	Z	m
Z
 d dlmZ d dlmZ d dlmZ d dlmZ d d	lmZmZ erGd d
lmZ neddd\ZZddgZdddddddddddejdejfd=d1dZ	2	3	4	4d>d?d<dZdS )@    )annotations)TYPE_CHECKINGAnyN)	DtypeLikeNdarrayOrTensor)
SpatialPad)rescale_array)repeat)optional_import)convert_data_typeconvert_to_dst_type)pyplot
matplotlibr   )name	matshow3dblend_images)
   r      noneFvolumer   figr   title
str | Nonefigsizetuple[int, int]frames_per_row
int | None	frame_dimintchannel_dimvminfloat | Nonevmaxevery_ninterpolationstrshowbool
fill_valuemargindtyper   kwargsreturntuple[Any, np.ndarray]c                   s  t | tjdd |dur|dvsj| dvrtdtttfrUtd tjs/tdtjt	dd	 D dd
}t
|dd tjfdd	D dd
ntjdk ritdtjdk s\|durt|ddjd jd jd fnt|ddjd jd f|du rtn|}|du rtn|}ddt|	d |stttt}ttt|d ttt  }d | t gg}|dur|ddgg7 }|||ggd 7 }tjj|dd|d|dt fdd	t|D }|dur't|dd}t|tjr1|}n|du r<tjdd}|jsE|d |jd }|j|f|||
d| | d |durf|!| |durvt"|drv|#| |r}t$  ||fS )a8
  
    Create a 3D volume figure as a grid of images.

    Args:
        volume: 3D volume to display. data shape can be `BCHWD`, `CHWD` or `HWD`.
            Higher dimensional arrays will be reshaped into (-1, H, W, [C]), `C` depends on `channel_dim` arg.
            A list of channel-first (C, H[, W, D]) arrays can also be passed in,
            in which case they will be displayed as a padded and stacked volume.
        fig: matplotlib figure or Axes to use. If None, a new figure will be created.
        title: title of the figure.
        figsize: size of the figure.
        frames_per_row: number of frames to display in each row. If None, sqrt(firstdim) will be used.
        frame_dim: for higher dimensional arrays, which dimension from (`-1`, `-2`, `-3`) is moved to
            the `-3` dimension. dim and reshape to (-1, H, W) shape to construct frames, default to `-3`.
        channel_dim: if not None, explicitly specify the channel dimension to be transposed to the
            last dimensionas shape (-1, H, W, C). this can be used to plot RGB color image.
            if None, the channel dimension will be flattened with `frame_dim` and `batch_dim` as shape (-1, H, W).
            note that it can only support 3D input image. default is None.
        vmin: `vmin` for the matplotlib `imshow`.
        vmax: `vmax` for the matplotlib `imshow`.
        every_n: factor to subsample the frames so that only every n-th frame is displayed.
        interpolation: interpolation to use for the matplotlib `matshow`.
        show: if True, show the figure.
        fill_value: value to use for the empty part of the grid.
        margin: margin to use for the grid.
        dtype: data type of the output stacked frames.
        kwargs: additional keyword arguments to matplotlib `matshow` and `imshow`.

    See Also:
        - https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.imshow.html
        - https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.matshow.html

    Example:

        >>> import numpy as np
        >>> import matplotlib.pyplot as plt
        >>> from monai.visualize import matshow3d
        # create a figure of a 3D volume
        >>> volume = np.random.rand(10, 10, 10)
        >>> fig = plt.figure()
        >>> matshow3d(volume, fig=fig, title="3D Volume")
        >>> plt.show()
        # create a figure of a list of channel-first 3D volumes
        >>> volumes = [np.random.rand(1, 10, 10, 10), np.random.rand(1, 10, 10, 10)]
        >>> fig = plt.figure()
        >>> matshow3d(volumes, fig=fig, title="List of Volumes")
        >>> plt.show()

    )dataoutput_typer   N)r   r   )r         zKchannel_dim must be: None, 0 or 1, and channels of image must be 1, 3 or 4.z volume must be a list of arrays.c                 S  s   g | ]}|j qS  )shape.0vr3   r3   W/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/visualize/utils.py
<listcomp>n   s    zmatshow3d.<locals>.<listcomp>axisr   c                   s   g | ]} |qS r3   r3   r5   )padr3   r8   r9   p   s    r1   r      F)copyconstant)modeconstant_valuesc                   s&   g | ]  fd dt D qS )c                   s   g | ]
}  |  qS r3   r3   )r6   j)colsivolr3   r8   r9      s    z(matshow3d.<locals>.<listcomp>.<listcomp>)range)r6   )rF   rH   )rG   r8   r9      s   & T)tight_layouto   )r!   r#   r%   offset_size_inches)%r   npndarrayr4   
ValueError
isinstancelisttuplemaxasarrayr   concatenatelenexpand_dimsmoveaxisreshapenanminnanmaxr   ceilsqrtminr<   astypeblockrI   pltAxesfigureaxesadd_subplotmatshowr;   	set_titlehasattrrM   r'   )r   r   r   r   r   r   r    r!   r#   r$   r%   r'   r)   r*   r+   r,   pad_sizerowswidthimaxr3   )rF   r<   rH   r8   r   "   sb   C&







      ?hsvTimagelabelalphafloat | NdarrayOrTensorcmaprescale_arraystransparent_backgroundc           
      C  s0  |j d dkrtd| j d dvrtd| j dd |j dd kr(tdt|tjtjfrC| j dd |j dd krCtd|rMt| } t|}| j d dkr[t| d	dd
} ddd}|||}t|tjtjfrq|}nt|tjr~t	||}nt	||}|rd||dk< d| }	|	|  ||  S )a  
    Blend an image and a label. Both should have the shape CHW[D].
    The image may have C==1 or 3 channels (greyscale or RGB).
    The label is expected to have C==1.

    Args:
        image: the input image to blend with label data.
        label: the input label to blend with image data.
        alpha: this specifies the weighting given to the label, where 0 is completely
            transparent and 1 is completely opaque. This can be given as either a
            single value or an array/tensor that is the same size as the input image.
        cmap: specify colormap in the matplotlib, default to `hsv`, for more details, please refer to:
            https://matplotlib.org/2.0.2/users/colormaps.html.
        rescale_arrays: whether to rescale the array to [0, 1] first, default to `True`.
        transparent_background: if true, any zeros in the label field will not be colored.

    .. image:: ../../docs/images/blend_images.png

    r   r   zLabel should have 1 channel.)r   r1   z"Image should have 1 or 3 channels.Nz3image and label should have matching spatial sizes.z;if alpha is image, size should match input image and label.r1   r:   ru   r&   rr   r   r-   c                 S  sP   t j| }t|tj^}}||d }t|ddd d }t||^}}|S )Nr   r>   r1   )rb   	colormapsget_cmapr   rN   rO   rY   r   )ru   rr   Z_cmapZlabel_np_Zlabel_rgb_np	label_rgbr3   r3   r8   get_label_rgb   s   z#blend_images.<locals>.get_label_rgb)ru   r&   rr   r   r-   r   )
r4   rP   rQ   rN   rO   torchTensorr   r	   	full_like)
rq   rr   rs   ru   rv   rw   r|   r{   Zw_labelZw_imager3   r3   r8   r      s2   

)"r   r   r   r   r   r   r   r   r   r   r   r   r    r   r!   r"   r#   r"   r$   r   r%   r&   r'   r(   r)   r   r*   r   r+   r   r,   r   r-   r.   )ro   rp   TT)rq   r   rr   r   rs   rt   ru   r&   rv   r(   rw   r(   r-   r   )
__future__r   typingr   r   numpyrN   r}   monai.config.type_definitionsr   r   monai.transforms.croppad.arrayr   monai.transforms.utilsr   0monai.transforms.utils_pytorch_numpy_unificationr	   monai.utils.moduler
   monai.utils.type_conversionr   r   r   r   rb   rz   __all__nanfloat32r   r   r3   r3   r3   r8   <module>   sF    