o eb@sXdZddlmZmZmZddlZddlmZ ddl m Z ddl mZGdd d eZdS) zCClasses to represent heatmaps, i.e. float arrays of ``[0.0, 1.0]``.)print_functiondivisionabsolute_importN)imgaug) IAugmentablec@seZdZdZd1ddZddZd2d d Zd3ddZddZd4ddZ  d5ddZ ddZ ddZ e jdddd d!Zd6d#d$Zd%d&Zed1d'd(Zed1d)d*Zed+d,Zd-d.Zd/d0ZdS)7HeatmapsOnImageaObject representing heatmaps on a single image. Parameters ---------- arr : (H,W) ndarray or (H,W,C) ndarray Array representing the heatmap(s) on a single image. Multiple heatmaps may be provided, in which case ``C`` is expected to denote the heatmap index. The array must be of dtype ``float32``. shape : tuple of int Shape of the image on which the heatmap(s) is/are placed. **Not** the shape of the heatmap(s) array, unless it is identical to the image shape (note the likely difference between the arrays in the number of channels). This is expected to be ``(H, W)`` or ``(H, W, C)`` with ``C`` usually being ``3``. If there is no corresponding image, use ``(H_arr, W_arr)`` instead, where ``H_arr`` is the height of the heatmap(s) array (analogous ``W_arr``). min_value : float, optional Minimum value for the heatmaps that `arr` represents. This will usually be ``0.0``. max_value : float, optional Maximum value for the heatmaps that `arr` represents. This will usually be ``1.0``. ?c Cst|sJdt|f|jddkr|jddks$Jd|jf|jjdvs2Jd|jf|jdvs?Jd|jft|dvsLJd |f||ksXJd ||ft |jj }|j dd }t |||k}t |||k}|s||rtd ||t |t |ft|||}|jd kr|dtjf}d|_nd|_d||kod|kn} d||kod|kn} | r| r||_n |||||_||_||_||_dS)z'Construct a new HeatmapsOnImage object.z8Expected numpy array as heatmap input array, got type %srrzSExpected numpy array as heatmap with height and width greater than 0, got shape %s.)float32zBHeatmap input array expected to be of dtype float32, got dtype %s.)rz3Heatmap input array must be 2d or 3d, got shape %s.zJArgument 'shape' in HeatmapsOnImage expected to be 2d or 3d, got shape %s.z@Expected min_value to be lower than max_value, got %.4f and %.4f2zValue range of heatmap was chosen to be (%.8f, %.8f), but found actual min/max of (%.8f, %.8f). Array will be clipped to chosen value range.r.TFr r N)ia is_np_arraytypeshapedtypenamendimlennpfinfoepsZflatminmaxwarnclipnewaxis arr_was_2darr_0to1 min_value max_value) selfarrrr!r"r componentsZ beyond_minZ beyond_max min_is_zero max_is_oner(LD:\Projects\ConvertPro\env\Lib\site-packages\imgaug/augmentables/heatmaps.py__init__+sl       zHeatmapsOnImage.__init__cCs|jr|jjddkr|jdddddf}n|j}ttjj}d||jko0d|kn}d||jkoAd|kn}|rM|rMt |S|j|j}|j||S)aGet the heatmap's array in value range provided to ``__init__()``. The :class:`HeatmapsOnImage` object saves heatmaps internally in the value range ``[0.0, 1.0]``. This function converts the internal representation to ``[min, max]``, where ``min`` and ``max`` are provided to :func:`HeatmapsOnImage.__init__` upon instantiation of the object. Returns ------- (H,W) ndarray or (H,W,C) ndarray Heatmap array of dtype ``float32``. rrNrr r ) rr rrrr rr!r"copy)r#r$rr&r'diffr(r(r)get_arrbs""  zHeatmapsOnImage.get_arrNjetc Cs|}g}t|jdD]^}|d||df}|dur'tj||dd}n|}t|tj d}|durOddl m }| |} | |} t | d d} n t|dtjfd } t| d dd tj} || q|S) aRender the heatmaps as RGB images. Parameters ---------- size : None or float or iterable of int or iterable of float, optional Size of the rendered RGB image as ``(height, width)``. See :func:`~imgaug.imgaug.imresize_single_image` for details. If set to ``None``, no resizing is performed and the size of the heatmaps array is used. cmap : str or None, optional Name of the ``matplotlib`` color map to use when convert the heatmaps to RGB images. If set to ``None``, no color map will be used and the heatmaps will be converted to simple intensity maps. Returns ------- list of (H,W,3) ndarray Rendered heatmaps as ``uint8`` arrays. Always a **list** containing one RGB image per heatmap array channel. r.rNZnearest interpolationo@rr )rrr )to_uint8smxrangerrimresize_single_imagerZsqueezeastyper Zmatplotlib.pyplotZpyplotZget_cmapdeleteZtilerruint8append) r#sizecmapZheatmaps_uint8heatmaps_drawncZ heatmap_cZ heatmap_c_rsZpltZ cmap_funcZheatmap_cmappedr(r(r)draws2    zHeatmapsOnImage.draw?heatmapscsjdksJdjjfjddks Jdjdfjjdks/Jdjjfdkr9dks@nJd |d vsKJd |f|d kr]tj|jjd ddd|j|dkrjjd dnd|d}fdd|D}|S)a7Draw the heatmaps as overlays over an image. Parameters ---------- image : (H,W,3) ndarray Image onto which to draw the heatmaps. Expected to be of dtype ``uint8``. alpha : float, optional Alpha/opacity value to use for the mixing of image and heatmaps. Larger values mean that the heatmaps will be more visible and the image less visible. cmap : str or None, optional Name of the ``matplotlib`` color map to use. See :func:`HeatmapsOnImage.draw` for details. resize : {'heatmaps', 'image'}, optional In case of size differences between the image and heatmaps, either the image or the heatmaps can be resized. This parameter controls which of the two will be resized to the other's size. Returns ------- list of (H,W,3) ndarray Rendered overlays as ``uint8`` arrays. Always a **list** containing one RGB image per heatmap array channel. r zUExpected to draw on three-dimensional image, got %d dimensions with shape %s instead.rz,Expected RGB image, got %d channels instead.r9z#Expected uint8 image, got dtype %s.g:0yEg1?z;Expected 'alpha' to be in the interval [0.0, 1.0], got %.4f)rAimagezsz1HeatmapsOnImage.draw_on_image..)rrrrrr6r r?)r#rBrFr<resizer=Zmixr(rEr) draw_on_imagesD    zHeatmapsOnImage.draw_on_imagecCs*tjd|j|j|j|jd}|j|_|S)aInvert each component in the heatmap. This shifts low values towards high values and vice versa. This changes each value to:: v' = max - (v - min) where ``v`` is the value at a spatial location, ``min`` is the minimum value in the heatmap and ``max`` is the maximum value. As the heatmap uses internally a ``0.0`` to ``1.0`` representation, this simply becomes ``v' = 1.0 - v``. This function can be useful e.g. when working with depth maps, where algorithms might have an easier time representing the furthest away points with zeros, requiring an inverted depth map. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Inverted heatmap. rrr!r")r from_0to1r rr!r"r)r#Zarr_invr(r(r)invertszHeatmapsOnImage.invertrconstantc Cs>ddlm}|j|j||||||d}tj||j|j|jdS)aPPad the heatmaps at their top/right/bottom/left side. Parameters ---------- top : int, optional Amount of pixels to add at the top side of the heatmaps. Must be ``0`` or greater. right : int, optional Amount of pixels to add at the right side of the heatmaps. Must be ``0`` or greater. bottom : int, optional Amount of pixels to add at the bottom side of the heatmaps. Must be ``0`` or greater. left : int, optional Amount of pixels to add at the left side of the heatmaps. Must be ``0`` or greater. mode : string, optional Padding mode to use. See :func:`~imgaug.imgaug.pad` for details. cval : number, optional Value to use for padding `mode` is ``constant``. See :func:`~imgaug.imgaug.pad` for details. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Padded heatmaps of height ``H'=H+top+bottom`` and width ``W'=W+left+right``. rr;)toprightbottomleftmodecvalrJ) augmentersr;padr r rKrr!r") r#rOrPrQrRrSrTiasizearr_0to1_paddedr(r(r)rV s # zHeatmapsOnImage.padFc CsNddlm}|j|j|||dd\}}tj||j|j|jd}|r%||fS|S)aZPad the heatmaps until they match a target aspect ratio. Depending on which dimension is smaller (height or width), only the corresponding sides (left/right or top/bottom) will be padded. In each case, both of the sides will be padded equally. Parameters ---------- aspect_ratio : float Target aspect ratio, given as width/height. E.g. ``2.0`` denotes the image having twice as much width as height. mode : str, optional Padding mode to use. See :func:`~imgaug.imgaug.pad` for details. cval : number, optional Value to use for padding if `mode` is ``constant``. See :func:`~imgaug.imgaug.pad` for details. return_pad_amounts : bool, optional If ``False``, then only the padded instance will be returned. If ``True``, a tuple with two entries will be returned, where the first entry is the padded instance and the second entry are the amounts by which each array side was padded. These amounts are again a tuple of the form ``(top, right, bottom, left)``, with each value being an integer. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Padded heatmaps as :class:`HeatmapsOnImage` instance. tuple of int Amounts by which the instance's array was padded on each side, given as a tuple ``(top, right, bottom, left)``. This tuple is only returned if `return_pad_amounts` was set to ``True``. rrNT) aspect_ratiorSrTreturn_pad_amountsrJ) rUr;pad_to_aspect_ratior r rKrr!r") r#rYrSrTrZrWrXZ pad_amountsrAr(r(r)r[Ss" * z#HeatmapsOnImage.pad_to_aspect_ratiocCs*tj|j|dd}tj||j|j|jdS)aAverage-pool the heatmap(s) array using a given block/kernel size. Parameters ---------- block_size : int or tuple of int Size of each block of values to pool, aka kernel size. See :func:`~imgaug.imgaug.pool` for details. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Heatmaps after average pooling. r )Zpad_cvalrJ)ravg_poolr r rKrr!r"r# block_sizeZarr_0to1_reducedr(r(r)r\szHeatmapsOnImage.avg_poolcCs&t|j|}tj||j|j|jdS)aMax-pool the heatmap(s) array using a given block/kernel size. Parameters ---------- block_size : int or tuple of int Size of each block of values to pool, aka kernel size. See :func:`~imgaug.imgaug.pool` for details. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Heatmaps after max-pooling. rJ)rmax_poolr r rKrr!r"r]r(r(r)r_szHeatmapsOnImage.max_poolzHeatmapsOnImage.resize()z(resize() has the exactly same interface.)Zalt_funccommentcOs|j|i|S)zBResize the heatmap(s) array given a target size and interpolation.)rH)r#argskwargsr(r(r)scaleszHeatmapsOnImage.scalerCcCs8tj|j||d}t|dd}tj||j|j|j dS)adResize the heatmap(s) array given a target size and interpolation. Parameters ---------- sizes : float or iterable of int or iterable of float New size of the array in ``(height, width)``. See :func:`~imgaug.imgaug.imresize_single_image` for details. interpolation : None or str or int, optional The interpolation to use during resize. See :func:`~imgaug.imgaug.imresize_single_image` for details. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Resized heatmaps object. r/r r rJ) rr6r rrr rKrr!r")r#sizesr0Zarr_0to1_resizedr(r(r)rHszHeatmapsOnImage.resizecCs*tt|jddd}|tj}|S)zConvert this heatmaps object to an ``uint8`` array. Returns ------- (H,W,C) ndarray Heatmap as an ``uint8`` array, i.e. with the discrete value range ``[0, 255]``. r2r)rrroundr r7r9)r#Z arr_0to255 arr_uint8r(r(r)r3s zHeatmapsOnImage.to_uint8cCs"|tjd}tj||||dS)aCreate a ``float``-based heatmaps object from an ``uint8`` array. Parameters ---------- arr_uint8 : (H,W) ndarray or (H,W,C) ndarray Heatmap(s) array, where ``H`` is height, ``W`` is width and ``C`` is the number of heatmap channels. Expected dtype is ``uint8``. shape : tuple of int Shape of the image on which the heatmap(s) is/are placed. **Not** the shape of the heatmap(s) array, unless it is identical to the image shape (note the likely difference between the arrays in the number of channels). If there is not a corresponding image, use the shape of the heatmaps array. min_value : float, optional Minimum value of the float heatmaps that the input array represents. This will usually be 0.0. In most other cases it will be close to the interval ``[0.0, 1.0]``. Calling :func:`~imgaug.HeatmapsOnImage.get_arr`, will automatically convert the interval ``[0.0, 1.0]`` float array to this ``[min, max]`` interval. max_value : float, optional Minimum value of the float heatmaps that the input array represents. This will usually be 1.0. See parameter `min_value` for details. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Heatmaps object. r1r!r")r7rr r rK)rfrr!r"r r(r(r) from_uint8s &zHeatmapsOnImage.from_uint8cCs t||ddd}||_||_|S)aCreate a heatmaps object from a ``[0.0, 1.0]`` float array. Parameters ---------- arr_0to1 : (H,W) or (H,W,C) ndarray Heatmap(s) array, where ``H`` is the height, ``W`` is the width and ``C`` is the number of heatmap channels. Expected dtype is ``float32``. shape : tuple of ints Shape of the image on which the heatmap(s) is/are placed. **Not** the shape of the heatmap(s) array, unless it is identical to the image shape (note the likely difference between the arrays in the number of channels). If there is not a corresponding image, use the shape of the heatmaps array. min_value : float, optional Minimum value of the float heatmaps that the input array represents. This will usually be 0.0. In most other cases it will be close to the interval ``[0.0, 1.0]``. Calling :func:`~imgaug.HeatmapsOnImage.get_arr`, will automatically convert the interval ``[0.0, 1.0]`` float array to this ``[min, max]`` interval. max_value : float, optional Minimum value of the float heatmaps that the input array represents. This will usually be 1.0. See parameter `min_value` for details. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Heatmaps object. r r rg)r r!r")r rr!r"rAr(r(r)rKs &zHeatmapsOnImage.from_0to1cCs,t|sJdt|fdd}t|tr|j|jf}n|d|t|tr/|j|jf}n|d|t|j j }|dd||dkoR|dd|kn}|dd||dkol|dd|kn}|rx|rxt |S|\}} |\} } | |} | | } ||| }| || }|S) aChange the value range of a heatmap array. E.g. the value range may be changed from the interval ``[0.0, 1.0]`` to ``[-1.0, 1.0]``. Parameters ---------- arr : ndarray Heatmap array to modify. source : tuple of float Current value range of the input array, given as a tuple ``(min, max)``, where both are ``float`` values. target : tuple of float Desired output value range of the array, given as a tuple ``(min, max)``, where both are ``float`` values. Returns ------- ndarray Input array, with value range projected to the desired target value range. z-Expected 'arr' to be an ndarray, got type %s.cSslt|tsJd|t|ft|dksJd|t|f|d|dks4Jd||d|dfdS)NzO'%s' was not a HeatmapsOnImage instance, expected type tuple then. Got type %s.rz;Expected tuple '%s' to contain exactly two entries, got %d.rrzYExpected tuple '%s' to have two entries with entry 1 < entry 2, got values %.4f and %.4f.) isinstancetuplerr)Zarg_name arg_valuer(r(r)_validate_tupleis"  z=HeatmapsOnImage.change_normalization.._validate_tuplesourcetargetr r) rrrrir r!r"rrrrr+)clsr$rmrnrlrZ mins_sameZ maxs_sameZ min_sourceZ max_sourceZ min_targetZ max_targetZ diff_sourceZ diff_targetr Z arr_targetr(r(r)change_normalizationKs,      44   z$HeatmapsOnImage.change_normalizationcCs|S)zCreate a shallow copy of the heatmaps object. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Shallow copy. )deepcopyr#r(r(r)r+s zHeatmapsOnImage.copycCst||j|j|jdS)zCreate a deep copy of the heatmaps object. Returns ------- imgaug.augmentables.heatmaps.HeatmapsOnImage Deep copy. rJ)r r-rr!r"rsr(r(r)rrs  zHeatmapsOnImage.deepcopy)r r )Nr.)r@r.rA)rrrrrMr )rMr F)rC)__name__ __module__ __qualname____doc__r*r-r?rIrLrVr[r\r_r deprecatedrcrHr3 staticmethodrhrK classmethodrqr+rrr(r(r(r)r s6 7 ! :C 3 ;  ! + , J r )rw __future__rrrnumpyrZ six.movesmovesr4rrbaserr r(r(r(r)s