o / iGT@sddlmZddlZddlmZmZddlmZddlm Z ddl m Z ddl m Z mZddlZddlZddlmZmZmZmZmZmZe rKddlZned \ZZgd ZGd d d eZGd ddeZed dGdddeZdddZdS)) annotationsN)ABCabstractmethod)Sequence) nullcontext)TemporaryDirectory) TYPE_CHECKINGAny)deprecated_argensure_tuple_sizeget_package_versionoptional_import require_pkg version_geqzarr)Merger AvgMerger ZarrAvgMergerc@s<eZdZdZ  ddd d ZedddZedddZdS)rao A base class for merging patches. Extend this class to support operations for `PatchInference`. There are two methods that must be implemented in the concrete classes: - aggregate: aggregate the values at their corresponding locations - finalize: perform any final process and return the merged output Args: merged_shape: the shape of the tensor required to merge the patches. cropped_shape: the shape of the final merged output tensor. If not provided, it will be the same as `merged_shape`. device: the device where Merger tensors should reside. N merged_shape Sequence[int] cropped_shapeSequence[int] | Nonedevicetorch.device | str | NonereturnNonecCsB|durtdt||_t|dur|jn||_||_d|_dS)Nz(Argument `merged_shape` must be providedF) ValueErrortuplerrr is_finalized)selfrrrr W/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/inferers/merger.py__init__9s   zMerger.__init__values torch.TensorlocationcCtd|jjd)a Aggregate values for merging. This method is being called in a loop and should add values to their corresponding location in the merged output results. Args: values: a tensor of shape BCHW[D], representing the values of inference output. location: a tuple/list giving the top left location of the patch in the output. Raises: NotImplementedError: When the subclass does not override this method. Subclass  must implement this method.NotImplementedError __class____name__)rr#r%r r r! aggregateGszMerger.aggregater cCr&)a Perform final operations for merging patches and return the final merged output. Returns: The results of merged patches, which is commonly a torch.Tensor representing the merged result, or a string representing the filepath to the merged results on disk. Raises: NotImplementedError: When the subclass does not override this method. r'r(r)rr r r!finalizeWs zMerger.finalize)NN)rrrrrrrrr#r$r%rrr)rr )r, __module__ __qualname____doc__r"rr-r/r r r r!r)s  rcsbeZdZdZdejejdfdfdd ZdddZd ddZ d ddZ d ddZ d ddZ Z S)!raMerge patches by taking average of the overlapping area Args: merged_shape: the shape of the tensor required to merge the patches. cropped_shape: the shape of the final merged output tensor. If not provided, it will be the same as `merged_shape`. device: the device for aggregator tensors and final results. value_dtype: the dtype for value aggregating tensor and the final result. count_dtype: the dtype for sample counting tensor. Ncpurrrr value_dtype torch.dtype count_dtypertorch.device | strrrcsjtj|||d|jstd|jd||_||_tj|j|j|jd|_ tj|j|j|jd|_ dS)N)rrrz1`merged_shape` must be provided for `AvgMerger`. is give.)dtyper) superr"rrr5r7torchzerosrr#counts)rrrr5r7rr+r r!r"sszAvgMerger.__init__r#r$r%cCsr|jrtd|jdd}tddt||D}t||jtddd}|j||7<|j |d7<dS) ab Aggregate values for merging. Args: values: a tensor of shape BCHW[D], representing the values of inference output. location: a tuple/list giving the top left location of the patch in the original image. Raises: NotImplementedError: When the subclass does not override this method. zO`AvgMerger` is already finalized. Please instantiate a new object to aggregate.Ncs"|] \}}t|||VqdSNslice.0locsizer r r!  z&AvgMerger.aggregate..Tpad_valpad_from_start) rrshaperzipr ndimrDr#r>rr#r% patch_sizeZ map_slicer r r!r-s zAvgMerger.aggregatecCs<|js|j|j|jtdd|jD|_d|_|jS)a Finalize merging by dividing values by counts and return the merged tensor. Notes: To avoid creating a new tensor for the final results (to save memory space), after this method is called, `get_values()` method will return the "final" averaged values, and not the accumulating values. Also calling `finalize()` multiple times does not have any effect. Returns: torch.tensor: a tensor of merged patches css|]}td|VqdS)rNrC)rFendr r r!rIsz%AvgMerger.finalize..T)rr#div_r>rrr.r r r!r/s  zAvgMerger.finalizecCs|S)zi Get the final merged output. Returns: torch.Tensor: merged output. )r/r.r r r! get_outputszAvgMerger.get_outputcC|jS)a Get the accumulated values during aggregation or final averaged values after it is finalized. Returns: torch.tensor: aggregated values. Notes: - If called before calling `finalize()`, this method returns the accumulating values. - If called after calling `finalize()`, this method returns the final merged [and averaged] values. r#r.r r r! get_valuess zAvgMerger.get_valuescCrW)z Get the aggregator tensor for number of samples. Returns: torch.Tensor: number of accumulated samples at each location. r>r.r r r! get_countszAvgMerger.get_counts) rrrrr5r6r7r6rr8rrr0)rr$)r,r1r2r3r<float32uint8r"r-r/rVrYr[ __classcell__r r r?r!rgs     r)pkg_namec seZdZdZeddddddedddd d ded ddd d d               d8d9fd(d) Zd:d-d.Zd;d0d1Zd;d2d3Zd;d4d5Z d;d6d7Z Z S) threadingLocklockr)rrrr:r5r7rnrprrrarhrjrbrirkrurwZ is_zarr_v3r?r r!r"s$           zZarrAvgMerger.__init__r#r$r%cCs|jrtd|jdd}tddt||D}t||jtddd}|j|j || 7<|j |d7<WddS1sHwYdS) z Aggregate values for merging. Args: values: a tensor of shape BCHW[D], representing the values of inference output. location: a tuple/list giving the top left location of the patch in the original image. zS`ZarrAvgMerger` is already finalized. Please instantiate a new object to aggregate.r@NcsrArBrCrEr r r!rIrJz*ZarrAvgMerger.aggregate..TrKrN) rrrOrrPr rQrDrr#numpyr>rRr r r!r-s"zZarrAvgMerger.aggregate zarr.ArraycCsR|js&t|jj|jjD]}|j||j||j|<q |j|jd|_|jS)a Finalize merging by dividing values by counts and return the merged tensor. Notes: To avoid creating a new tensor for the final results (to save memory space), after this method is called, `get_values()` method will return the "final" averaged values, and not the accumulating values. Also calling `finalize()` multiple times does not have any effect. Returns: zarr.Array: a zarr array of of merged patches T) riterate_over_chunksr#ru cdata_shaper>rresizer)rchunkr r r!r/s  zZarrAvgMerger.finalizecCrW)zy Get the final merged output. Returns: zarr.Array: Merged (averaged) output tensor. )rr.r r r!rVr\zZarrAvgMerger.get_outputcCrW)z} Get the accumulated values during aggregation Returns: zarr.Array: aggregated values. rXr.r r r!rYszZarrAvgMerger.get_valuescCrW)z Get the aggregator tensor for number of samples. Returns: zarr.Array: Number of accumulated samples at each location. rZr.r r r!r[r\zZarrAvgMerger.get_counts)Nr]r]r^rlNNNNNNNNTT)"rrrrr:rmr5rmr7rmrnrorprqrrrqrarsrhrsrjrsrbrtrirtrkrtrurvrwrxrrr0)rr) r,r1r2r3r r"r-r/rVrYr[r_r r r?r!rsR&   (    rr c cst|t|kr tdt|dkr1t|dD]}|t||d|d|dfVqdSt|dD]#}t|dd|dd|t||d|d|dfEdHq7dS)aQ Iterate over chunks of a given shape. Args: chunks: the chunk shape cdata_shape: the shape of the data in chunks slice_tuple: the slice tuple to be used for indexing Raises: ValueError: When the length of chunks and cdata_shape are not the same. Yields: slices of the data z0chunks and cdata_shape must have the same lengthrNrN)lenrrangerDr)rurZ slice_tupleir r r!rs (6 r)r ) __future__rrabcrrcollections.abcr contextlibrtempfilertypingrr rnpr< monai.utilsr r r r rrr___all__rrrrr r r r!s*       >g'