o - i @sfddlmZddlZddlmZddlmZddlmZeddd\Z Z d gZ Gd d d ejj Z dS) ) annotationsN)nn) functional)optional_importzsegment_anything.build_sambuild_sam_vit_b)nameCellSamWrappercs4eZdZdZ    d d fdd Zd d ZZS)ra) CellSamWrapper is thin wrapper around SAM model https://github.com/facebookresearch/segment-anything with an image only decoder, that can be used for segmentation tasks. Args: auto_resize_inputs: whether to resize inputs before passing to the network. (usually they need be resized, unless they are already at the expected size) network_resize_roi: expected input size for the network. (currently SAM expects 1024x1024) checkpoint: checkpoint file to load the SAM weights from. (this can be downloaded from SAM repo https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth) return_features: whether to return features from SAM encoder (without using decoder/upsampling to the original input size) Tr sam_vit_b_01ec64.pthFreturnNonecstj|i|||_||_||_tstdt|d}d|_d|_ t t j ddt j ddt jdddd d d d d t j ddt j ddt jdddd d d dd |_ ||_dS) NzjSAM is not installed, please run: pip install git+https://github.com/facebookresearch/segment-anything.git) checkpoint) num_featuresT)inplaceF) kernel_sizestridepaddingoutput_paddingbias)super__init__network_resize_roiauto_resize_inputsreturn_featureshas_sam ValueErrorrZprompt_encoder mask_decoderr Sequential BatchNorm2dReLUConvTranspose2dmodel)selfrrrrargskwargsr' __class__f/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/networks/nets/cell_sam_wrapper.pyr+s(       zCellSamWrapper.__init__cCs^|jdd}|jrtj||jdd}|j|}|js-|j|}|jr-tj||dd}|S)Nrbilinear)sizemode) shaperF interpolaterr' image_encoderrr")r(xshr-r-r.forwardOs  zCellSamWrapper.forward)Tr r F)r r )__name__ __module__ __qualname____doc__rr8 __classcell__r-r-r+r.rs$) __future__rtorchrtorch.nnrr3 monai.utilsrrr Z_all__Modulerr-r-r-r.s