o  i#=@sjdZddlmZddlZddlmZmZddlmZddl m Z m Z ddl m Z mZddlZddlZddlmZdd lmZmZz ddlmZd ZWn eyUd ZYnwe radd lmZmZned ejed\ZZ ed ejed\ZZ dZ!de!fdfdBd d!Z"de!fdfdCd'd(Z#dDd-d.Z$d/de!fde$ddfdEd8d9Z%dFdd?d?eZ'Gd@dAdAe Z(dS)Gz This set of utility function is meant to make using Jupyter notebooks easier with MONAI. Plotting functions using Matplotlib produce common plots for metrics and images. ) annotationsN)CallableMapping)Enum)RLockThread) TYPE_CHECKINGAny) IgniteInfo) min_versionoptional_importTF)EngineEventsz ignite.enginer rlosslogaxplt.Axestitlestrgraphmap;Mapping[str, list[float] | tuple[list[float], list[float]]]yscaleavg_keys tuple[str]window_fractionintreturnNonecCsHddlm}|D]p\}}t|dkrzt|dttfr$t|\} } n ttt|t|} } |j | | |d| ddd||vrzt||krzt||} t | f| } t j | df| d| | dd } |j | | |d | dddq | ||||d |jd dd d|ddd|j|dddS)ar Plot metrics on a single graph with running averages plotted for selected keys. The values in `graphmap` should be lists of (timepoint, value) pairs as stored in MetricLogger objects. Args: ax: Axes object to plot into title: graph title graphmap: dictionary of named graph values, which are lists of values or (index, value) pairs yscale: scale for y-axis compatible with `Axes.set_yscale` avg_keys: tuple of keys in `graphmap` to provide running average plots for window_fraction: what fraction of the graph value length to use as the running average window r) MaxNLocatorz = z.5g)labelvalid)modez Avg = on)r"r"g)Zbbox_to_anchorlocZ borderaxespadTboth)integerN)Zmatplotlib.tickerritemslen isinstancetuplelistziprangeplotnponesconvolve set_titleZ set_yscaleaxislegendgridZxaxisZset_major_locator)rrrrrrrnvindsvalswindowkernelrar?[/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/utils/jupyter_utils.pyplot_metric_graph.s&    "    rAfig plt.Figureimagemapdict[str, np.ndarray]r-c Csdtdt|f}tj|d|d|d}t|||||||g} t|D]R\} } tj|d| fd|d} || jddkrI| || gd n | jt || d d | | d ||  d d|| d | d| | q$| S)a Plot metric graph data with images below into figure `fig`. The intended use is for the graph data to be metrics from a training run and the images to be the batch and output from the last iteration. This uses `plot_metric_graph` to plot the metric graph. Args: fig: Figure object to plot into, reuse from previous plotting for flicker-free refreshing title: graph title graphmap: dictionary of named graph values, which are lists of values or (index, value) pairs imagemap: dictionary of named images to show with metric plot yscale: for metric plot, scale for y-axis compatible with `Axes.set_yscale` avg_keys: for metric plot, tuple of keys in `graphmap` to provide running average plots for window_fraction: for metric plot, what fraction of the graph value length to use as the running average window Returns: list of Axes objects for graph followed by images r")rr)colspanrB)rowspanrBr)r"rHrZgray)Zcmap z.3gz -> off)maxr*pltZ subplot2gridrA enumerateshapeZimshow transposer1squeezer4minr5append) rBrrrDrrrZ gridshapegraphaxesir8imr?r?r@plot_metric_images]s.  rYnametensor torch.Tensornp.ndarray | NonecCs|jdkr|jddkr|jddkr|jS|jdkrA|jddkrA|jddkrA|jdd}|dd|fjSdS)a Return an tuple of images derived from the given tensor. The `name` value indices which key from the output or batch value the tensor was stored as, or is "Batch" or "Output" if these were single tensors instead of dictionaries. Returns a tuple of 2D images of shape HW, or 3D images of shape CHW where C is color channels RGB or RGBA. This allows multiple images to be created from a single tensor, ie. to show each channel separately. rJr"rHrFN)ndimrPcpudatanumpy)rZr[Zdmidr?r?r@tensor_to_imagess &&rbz Training Logengineloggerr image_fn)Callable[[str, torch.Tensor], Any] | Noneplt.Figure | None selected_insttuple[plt.Figure, list]c Cs|dur |ntjdddd}t|ji} | |ji} |dur|jdur|jjdur|jj|jj fD]q} | |jjur?dnd} |} t | t rW| |}d} d d | D} t | t r| D],\}}t |tjr|jd kr|||| }|durt|D] \}}|| |d |<q~q`q5t | tjr|| | }|dur|| | d |<q5t||| | |||}|jr|dj|jd dddd||fS)a Plot the status of the given Engine with its logger. The plot will consist of a graph of loss values and metrics taken from the logger, and images taken from the `output` and `batch` members of `engine.state`. The images are converted to Numpy arrays suitable for input to `Axes.imshow` using `image_fn`, if this is None then no image plotting is done. Args: engine: Engine to extract images from logger: MetricLogger to extract loss and metric data from title: graph title yscale: for metric plot, scale for y-axis compatible with `Axes.set_yscale` avg_keys: for metric plot, tuple of keys in `graphmap` to provide running average plots for window_fraction: for metric plot, what fraction of the graph value length to use as the running average window image_fn: callable converting tensors keyed to a name in the Engine to a tuple of images to plot fig: Figure object to plot into, reuse from previous plotting for flicker-free refreshing selected_inst: index of the instance to show in the image plot Returns: Figure object (or `fig` if given), list of Axes objects for graph and images N)r TZwhite)figsizeZ tight_layoutZ facecolorZBatchOutputrcSs0i|]\}}t|tjr|jdkr||dqS)rJN)r+torchTensorr^).0kr9r?r?r@ s0z&plot_engine_status..rF_r r"rp:)cls)ZclfrNFigure LOSS_NAMErupdatemetricsstatebatchoutputr+r-r)dictrmrnr^rOrYZaxhline)rcrdrrrrrerBrhrrDsrcr!Zbatch_selected_instZ selected_dictrpr9imagerWrXrVr?r?r@plot_engine_statuss@      rr|Ulist[torch.Tensor | dict[str, torch.Tensor]] | dict[str, torch.Tensor] | torch.TensorcCs(d dd}t|tr||dS||S) zJReturns a single value from the network output, which is a dict or tensor.r`&torch.Tensor | dict[str, torch.Tensor]rr\cSst|tr |dS|S)Nr)r+r})r`r?r?r@ _get_losss z(_get_loss_from_output.._get_lossrN)r`rrr\)r+r-)r|rr?r?r@_get_loss_from_outputs   rc@s eZdZdZdZdZdZdZdS) StatusMembersz` Named members of the status dictionary, others may be present for named metric values. StatusZEpochsZItersZLossN)__name__ __module__ __qualname____doc__STATUSEPOCHSITERSLOSSr?r?r?r@rs rcsleZdZdZedddfd fd d ZddZddZddZe d!ddZ d"ddZ e fd#ddZ ZS)$ThreadContainera Contains a running `Engine` object within a separate thread from main thread in a Jupyter notebook. This allows an engine to begin a run in the background and allow the starting notebook cell to complete. A user can thus start a run and then navigate away from the notebook without concern for loosing connection with the running cell. All output is acquired through methods which synchronize with the running engine using an internal `lock` member, acquiring this lock allows the engine to be inspected while it's prevented from starting the next iteration. Args: engine: wrapped `Engine` object, when the container is started its `run` method is called loss_transform: callable to convert an output dict into a single numeric value metric_transform: callable to convert a named metric value into a single numeric value status_format: format string for status key-value pairs. cCs|SNr?)rZvaluer?r?r@szThreadContainer.z {}: {:.4}rcr loss_transformrmetric_transform status_formatrcsLtt|_||_i|_||_||_d|_||_ |j t j |j dSr)super__init__rlockrc _status_dictrrrBrZadd_event_handlerrZITERATION_COMPLETED_update_status)selfrcrrr __class__r?r@rs zThreadContainer.__init__cCs|jdS)z-Calls the `run` method of the wrapped engine.N)rcrunrr?r?r@r(szThreadContainer.runcCs|j|dS)z$Stop the engine and join the thread.N)rc terminatejoinrr?r?r@stop,s  zThreadContainer.stopcCs2|j|jj}tjjdtjjdtjjtdi}|dur|j dur0|j dkr0|j d|j }nt |j }|j durG|j |j d|j }nt |j }||tjj<||tjj<||j|tjj<|jpfi}|D]\}}|||}|dur|||qk|j|WddS1swYdS)zNCalled as an event, updates the internal status dict at the end of iterations.rnanNr"/)rrcrzrrrrrfloatZ max_epochsepochrZ epoch_length iterationrr|ryr)rrTrrx)rrzstatsritersrymr9r?r?r@r1s0        "zThreadContainer._update_statusrdict[str, str]cCsR|jtjj|r dndi}||j|WdS1s"wYdS)zTA dictionary containing status information, current loss, and current metric values.ZRunningZStoppedN)rrrris_aliverxr)rrr?r?r@ status_dictRs  $zThreadContainer.status_dictcCst|j}|tjjdt|tjjdg}| D]\}}t |t r.|j ||}n|d|}||qd|S)zsT      4 0 N