a hY@sdZddlmZddlmZmZddlZddlm Z ddl m Z ddl m Z dd lmZmZmZdd lmZeeZerddlZed d Gd dde ZdgZdS)z Processor class for SAM2. deepcopy)OptionalUnionN) ImageInput)ProcessorMixin) BatchEncoding) TensorTypeis_torch_availablelogging)requires)torch)backendsc sLeZdZdZdgZdZd%eeedfdd Zd&e e ee e e e e e e jfee e e e ee jfee e e e e e jfee e e e e jfee eefed d d Zd'ed d dddZd(ddZd)ddZd*ddZddZddZd+e e jeje feeeeee dddZd,dd Zd-d#d$ZZS). Sam2Processora Constructs a SAM2 processor which wraps a SAM2 image processor and an 2D points & Bounding boxes processor into a single processor. [`Sam2Processor`] offers all the functionalities of [`Sam2ImageProcessorFast`] and [`Sam2VideoProcessor`]. See the docstring of [`~Sam2ImageProcessorFast.__call__`] and [`~Sam2VideoProcessor.__call__`] for more information. Args: image_processor (`Sam2ImageProcessorFast`): An instance of [`Sam2ImageProcessorFast`]. target_size (`int`, *optional*): The target size (target_size, target_size) to which the image will be resized. point_pad_value (`int`, *optional*, defaults to -10): The value used for padding input points. image_processorZSam2ImageProcessorFastN) target_sizepoint_pad_valuec s8tj|fi|||_|dur&|n |jjd|_dS)Nheight)super__init__rrsizer)selfrrrkwargs __class__d/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/sam2/processing_sam2.pyr:szSam2Processor.__init__)imagessegmentation_maps input_points input_labels input_boxesoriginal_sizesreturn_tensorsreturnc sN|dur"|j|f||d|} n:|durTt|tjrB|}td|i|d} ntd| d}|durt|dkrt|t|krtd|dus|dus|durJ|j |dd d d d } |j |d ddd} |j |d dddd } | dur| | dd } | dur| | dd }| dur6| | dd | dur\| dur\| |kr\td| durt| d krt fdd| Drtd| dur| | | d g}tj |tjd}|j||dd| d|i| dur| | |}tj |tjd}| d|i| durJtj | tjd}|j||dd| d|i| S)a This method uses [`Sam2ImageProcessorFast.__call__`] method to prepare image(s) for the model. It also prepares 2D points and bounding boxes for the model if they are provided. Args: images (`ImageInput`, *optional*): The image(s) to process. segmentation_maps (`ImageInput`, *optional*): The segmentation maps to process. input_points (`list[list[list[list[float]]]]`, `torch.Tensor`, *optional*): The points to add to the frame. input_labels (`list[list[list[int]]]`, `torch.Tensor`, *optional*): The labels for the points. input_boxes (`list[list[list[float]]]`, `torch.Tensor`, *optional*): The bounding boxes to add to the frame. original_sizes (`list[list[float]]`, `torch.Tensor`, *optional*): The original sizes of the images. return_tensors (`str` or `TensorType`, *optional*): The type of tensors to return. **kwargs: Additional keyword arguments to pass to the image processor. Returns: A [`BatchEncoding`] with the following fields: - `pixel_values` (`torch.Tensor`): The processed image(s). - `original_sizes` (`list[list[float]]`): The original sizes of the images. - `reshaped_input_sizes` (`torch.Tensor`): The reshaped input sizes of the images. - `labels` (`torch.Tensor`): The processed segmentation maps (if provided). - `input_points` (`torch.Tensor`): The processed points. - `input_labels` (`torch.Tensor`): The processed labels. - `input_boxes` (`torch.Tensor`): The processed bounding boxes. N)r r%r$)Z tensor_typez0Either images or original_sizes must be providedz{original_sizes must be of length 1 or len(images). If you are passing a single image, you must pass a single original_size.Zpointsz;[image level, object level, point level, point coordinates])expected_depth input_nameexpected_formatexpected_coord_sizerlabelsz([image level, object level, point level])r*r+r,Zboxesz)[image level, box level, box coordinates]zbInput points and labels have inconsistent dimensions. Please ensure they have the same dimensions.c3s|]}t|dkVqdSr'N)len).0Z img_boxesZboxes_max_dimsrr z)Sam2Processor.__call__..zInput boxes have inconsistent dimensions that would require padding, but boxes cannot be padded due to model limitations. Please ensure all images have the same number of boxes.)ZdtypeT)preserve_paddingr!r"is_bounding_boxr#)r isinstancerTensorcputolistr ValueErrorr0_validate_single_input_get_nested_dimensionsany_pad_nested_listtensorZfloat32_normalize_tensor_coordinatesupdateZint64)rrr r!r"r#r$r%rZencoding_image_processorZprocessed_pointsZprocessed_labelsZprocessed_boxesZpoints_max_dimsZlabels_max_dimsZ padded_pointsZ final_pointsZ padded_labelsZ final_labelsZ final_boxesrr2r__call__?s+  $        zSam2Processor.__call__Fz torch.Tensor)rcoordsr&c Csl|\}}||}}t|}|r0|ddd}|d|||d<|d|||d<|rh|dd}|S)a  Expects a numpy array of length 2 in the final dimension. Requires the original image size in (H, W) format. Args: target_size (`int`): The target size of the image. coords (`torch.Tensor`): The coordinates to be normalized. original_size (`tuple`): The original size of the image. is_bounding_box (`bool`, *optional*, defaults to `False`): Whether the coordinates are bounding boxes. r)).r).r'r()rfloatZreshape) rrrE original_sizer7Zold_hZold_wZnew_hZnew_wrrr_normalize_coordinatess   z$Sam2Processor._normalize_coordinatesrcs|dur dSt|tjrVdks2t|jdkr>|Sfdd|DSnt|tjrdks|t|jdkr|Sfdd|DSnRt|t rʈkr|Sfdd|DSn$t|t t fr|St dt |dS)aS Recursively convert various input formats (tensors, numpy arrays, lists) to nested lists. Args: data: Input data in any format expected_depth: Expected nesting depth current_depth: Current depth in recursion Returns: Nested list representation of the data Nr)csg|]}|dqSr'_convert_to_nested_listr1item current_depthr*rrr r4z9Sam2Processor._convert_to_nested_list..csg|]}|dqSrJrKrMrOrrrQr4csg|]}|dqSrJrKrMrOrrrQr4zUnsupported data type: )r8rr9r0shapenumpyr;npndarraylistintrGr<type)rdatar*rPrrOrrLs"     z%Sam2Processor._convert_to_nested_listcCs|dur g}t|ts|St|dkr6|t|nt|dt||d<t|dkr|D]`}t|tr\||}t|D]>\}}|dt|kr||q|t||d|||d<q|q\|S)a` Get the maximum dimensions at each level of nesting. Args: nested_list (`list`): Nested list structure. max_dims (`list`, *optional*): Current maximum dimensions (for recursion). Returns: `list`: A list of maximum dimensions for each nesting level. Nrr')r8rVr0appendmaxr> enumerate)r nested_listZmax_dimsrNZsub_dimsidimrrrr>s       z$Sam2Processor._get_nested_dimensionsc s\|dur|j}|t|kr|St|ts.|g}t|}||}|t|dkrd||g||n|dkr|t|dkr||dd}|||n|g||d|fddt||Dn8||dd}||||fddt|D|t|dkrXtt|D]2}t||tr$|||||d|||<q$|S)a Recursively pad a nested list to match target dimensions. Args: nested_list (`list`): Nested list to pad. target_dims (`list`): Target dimensions for each level. current_level (`int`, *optional*, defaults to 0): Current nesting level. pad_value (`int`, *optional*): Value to use for padding. Returns: `list`: The padded nested list. Nr'rr)csg|] }tqSrrr1_templaterrrQ[r4z2Sam2Processor._pad_nested_list..csg|] }tqSrrr`rbrrrQ`r4)rr0r8rVextend_create_empty_nested_structureranger@) rr]Z target_dimsZ current_level pad_valueZ current_sizerZ template_dimsr^rrbrr@,s0  "  zSam2Processor._pad_nested_listcs<tdkrgdSfddtdDSdS)a Create an empty nested structure with given dimensions filled with pad_value. Args: dims (`list`): The dimensions of the nested structure. pad_value (`int`): The value to fill the structure with. r'rcs g|]}ddqSr/)rer`dimsrgrrrrQwr4z@Sam2Processor._create_empty_nested_structure..N)r0rf)rrirgrrhrrejs z,Sam2Processor._create_empty_nested_structurecCsLt|tr,t|dkrdSd||dSt|tjtjfrHt|jSdS)z Get the nesting level of a list structure. Args: input_list (`list`): The list to get the nesting level of. rr') r8rVr0_get_nesting_levelrTrUrr9rR)rZ input_listrrrrjys   z Sam2Processor._get_nesting_level)rYr*r+r,r-r&c Cs|dur dSt|tjtjfr|j|krNtd|d|d|d|jd n8|dur|jd|krtd|d|d|jdd |||St|t r| |}||krtd|d |d |d|d |||SdS) a Validate a single input by ensuring proper nesting and raising an error if the input is not valid. Args: data (`torch.Tensor`, `np.ndarray`, or `list`): Input data to process. expected_depth (`int`): Expected nesting depth. input_name (`str`): Name of the input for error messages. expected_format (`str`): The expected format of the input. expected_coord_size (`int`, *optional*): Expected coordinate size (2 for points, 4 for boxes, None for labels). . NzInput z must be a tensor/array with z, dimensions. The expected nesting format is z. Got z dimensions.rFz as the last dimension, got .z must be a nested list with z( levels. The expected nesting format is z levels.) r8rr9rTrUndimr<rRrLrVrj)rrYr*r+r,r-rPrrrr=s(    z$Sam2Processor._validate_single_inputc Cs|r||jk}|jddd}tt|D]|}||jdkr(|t|krN||n|d}|j|j||||d} |r||} t| ||| ||||<q(| ||<q(dS)a Helper method to normalize coordinates in a tensor across multiple images. Args: tensor (`torch.Tensor`): Input tensor with coordinates. original_sizes (`list`): Original image sizes. is_bounding_box (`bool`, *optional*, defaults to `False`): Whether coordinates are bounding boxes. preserve_padding (`bool`, *optional*, defaults to `False`): Whether to preserve padding values (for points). rFT)r_Zkeepdimrr6N) rallrfr0rRrIrrwhereZ expand_as) rrAr$r7r5maskZ coord_maskZimg_idxrHZnormalized_coordsZimg_maskrrrrBs  z+Sam2Processor._normalize_tensor_coordinatesTc Ks |jj|||||||fi|S)a- Remove padding and upscale masks to the original image size. Args: masks (`Union[List[torch.Tensor], List[np.ndarray]]`): Batched masks from the mask_decoder in (batch_size, num_channels, height, width) format. original_sizes (`Union[torch.Tensor, List[Tuple[int,int]]]`): The original sizes of each image before it was resized to the model's expected input shape, in (height, width) format. mask_threshold (`float`, *optional*, defaults to 0.0): Threshold for binarization and post-processing operations. binarize (`bool`, *optional*, defaults to `True`): Whether to binarize the masks. max_hole_area (`float`, *optional*, defaults to 0.0): The maximum area of a hole to fill. max_sprinkle_area (`float`, *optional*, defaults to 0.0): The maximum area of a sprinkle to fill. apply_non_overlapping_constraints (`bool`, *optional*, defaults to `False`): Whether to apply non-overlapping constraints to the masks. Returns: (`torch.Tensor`): Batched masks in batch_size, num_channels, height, width) format, where (height, width) is given by original_size. )rpost_process_masks) rZmasksr$Zmask_thresholdZbinarizeZ max_hole_areaZmax_sprinkle_areaZ!apply_non_overlapping_constraintsrrrrrqs#z Sam2Processor.post_process_masks)Nr)NNNNNNN)F)r)N)rN)N)FF)rpTrprpF)__name__ __module__ __qualname____doc__ attributesZimage_processor_classrrWrrrrVrGrr9strr r rDrIrLr>r@rerjrTrUr=rBrq __classcell__rrrrr%s`    & % > 2 'r)rucopyrtypingrrrSrTZ image_utilsrZprocessing_utilsrZtokenization_utils_baser utilsr r r Zutils.import_utilsr Z get_loggerrrloggerrr__all__rrrrs"      k