o ĎiWf@sddlmZddlZddlZddlmZddlmZddl m Z ddl Z ddl mZddlmZddlmZdd lmZmZdd lmZdd lmZmZmZmZdd lmZmZmZdd l m!Z!m"Z"dZ#     dfdgddZ$dhd%d&Z%did)d*Z&Gd+d,d,Z'Gd-d.d.Z(         /  0 djdkd;d<Z)  0 = =   /  0 dldmdDdEZ*dndodLdMZ+dpdOdPZ,dqdSdTZ-drdVdWZ.dsdYdZZ/dtd[d\Z0dud_d`Z1dvdddeZ2dS)w) annotationsN)Sequence)Path)Literal)Image)TopologicalError)tqdm) BoundingBoxMask)logger)CocoCocoAnnotation CocoImagecreate_coco_dict)IMAGE_EXTENSIONS_LOSSLESSIMAGE_EXTENSIONS_LOSSYread_image_as_pil) load_json save_jsonT皙? image_heightint image_width slice_height int | None slice_widthauto_slice_resolution bool | Noneoverlap_height_ratio float | Noneoverlap_width_ratioreturnlist[list[int]]cCsg}d}} |r|rt||} t||} n|r$t||d\} } }}ntd||krd} } | |}| |krw| |} ||ksD| |krft|| }t||}td||}td||}|||||gn || | | |g| | } | |ks8|| } ||ks,|S)aGenerate bounding boxes for slicing an image into crops. The function calculates the coordinates for each slice based on the provided image dimensions, slice size, and overlap ratios. If slice size is not provided and auto_slice_resolution is True, the function will automatically determine appropriate slice parameters. Args: image_height (int): Height of the original image. image_width (int): Width of the original image. slice_height (int, optional): Height of each slice. Default None. slice_width (int, optional): Width of each slice. Default None. overlap_height_ratio (float, optional): Fractional overlap in height of each slice (e.g. an overlap of 0.2 for a slice of size 100 yields an overlap of 20 pixels). Default 0.2. overlap_width_ratio(float, optional): Fractional overlap in width of each slice (e.g. an overlap of 0.2 for a slice of size 100 yields an overlap of 20 pixels). Default 0.2. auto_slice_resolution (bool, optional): if not set slice parameters such as slice_height and slice_width, it enables automatically calculate these parameters from image resolution and orientation. Returns: List[List[int]]: List of 4 corner coordinates for each N slices. [ [slice_0_left, slice_0_top, slice_0_right, slice_0_bottom], ... [slice_N_left, slice_N_top, slice_N_right, slice_N_bottom] ] rheightwidthzECompute type is not auto and slice width and height are not provided.)rget_auto_slice_params ValueErrorminmaxappend)rrrrrrr! slice_bboxesy_maxy_min y_overlap x_overlapx_minx_maxxmaxymaxxminyminr7H/home/jeff/fluffinator/venv/lib/python3.10/site-packages/sahi/slicing.pyget_slice_bboxess2&    r9 annotationdict slice_bbox list[int]boolcCsd|d\}}}}||}||}||dkrdS||dkr dS||dkr(dS||dkr0dSdS)aqCheck whether annotation coordinates lie inside slice coordinates. Args: annotation (dict): Single annotation entry in COCO format. slice_bbox (List[int]): Generated from `get_slice_bboxes`. Format for each slice bbox: [x_min, y_min, x_max, y_max]. Returns: (bool): True if any annotation coordinate lies inside slice. bboxFrTr7)r:r<lefttopr&r%rightbottomr7r7r8annotation_inside_sliceZs     rGcoco_annotation_listlist[CocoAnnotation]cCsBg}|D]}t|j|r||}|j|j|kr||q|S)a1Slices and filters given list of CocoAnnotation objects with given 'slice_bbox' and 'min_area_ratio'. Args: coco_annotation_list (List[CocoAnnotation]) slice_bbox (List[int]): Generated from `get_slice_bboxes`. Format for each slice bbox: [x_min, y_min, x_max, y_max]. min_area_ratio (float): If the cropped annotation area to original annotation ratio is smaller than this value, the annotation is filtered out. Default 0.1. Returns: (List[CocoAnnotation]): Sliced annotations. )rGjsonget_sliced_coco_annotationarear+)rHr<min_area_ratiosliced_coco_annotation_listcoco_annotationsliced_coco_annotationr7r7r8process_coco_annotationsvs   rQc@seZdZddZdS) SlicedImagecCs||_||_||_dS)a image: np.array Sliced image. coco_image: CocoImage Coco styled image object that belong to sliced image. starting_pixel: list of list of int Starting pixel coordinates of the sliced image. Nimage coco_imagestarting_pixel)selfrTrUrVr7r7r8__init__s  zSlicedImage.__init__N)__name__ __module__ __qualname__rXr7r7r7r8rRs rRc@steZdZddddZdd d Zed d ZeddZedddZed ddZ ed ddZ ddZ ddZ dS)!SliceImageResultNoriginal_image_sizer= image_dir str | NonecCs$|d|_|d|_||_g|_dS)z image_dir: str Directory of the sliced image exports. original_image_size: list of int Size of the unsliced original image in [height, width] rrBN)original_image_heightoriginal_image_widthr^_sliced_image_list)rWr]r^r7r7r8rXs   zSliceImageResult.__init__ sliced_imagerRcCs"t|ts td|j|dS)Nz+sliced_image must be a SlicedImage instance) isinstancerR TypeErrorrbr+)rWrcr7r7r8add_sliced_images z!SliceImageResult.add_sliced_imagecCs|jSN)rbrWr7r7r8sliced_image_listsz"SliceImageResult.sliced_image_listcC g}|jD]}||jq|S)zXReturns sliced images. Returns: images: a list of np.array )rbr+rT)rWimagesrcr7r7r8rk zSliceImageResult.imagesr"list[CocoImage]cCrj)z}Returns CocoImage representation of SliceImageResult. Returns: coco_images: a list of CocoImage )rbr+rU)rW coco_imagesrcr7r7r8rnrlzSliceImageResult.coco_imagescCrj)zReturns a list of starting pixels for each slice. Returns: starting_pixels: a list of starting pixel coords [x,y] )rbr+rV)rWstarting_pixelsrcr7r7r8rorlz SliceImageResult.starting_pixelscCs"g}|jD] }||jjq|S)zxReturns a list of filenames for each slice. Returns: filenames: a list of filenames as str )rbr+rU file_name)rW filenamesrcr7r7r8rqs zSliceImageResult.filenamescsfddt|tjr|}t|tr|St|tr5|t\}}}fddt|||DSt|t t frEt |}t |St t |)Ncs&j|j|j|j|dS)N)rTrUrVfilename)rkrnrorq)irhr7r8_prepare_ith_dicts z7SliceImageResult.__getitem__.._prepare_ith_dictcsg|]}|qSr7r7).0rs)rtr7r8 sz0SliceImageResult.__getitem__..)rdnpndarraytolistrsliceindiceslenrangetuplelistmapNotImplementedErrortype)rWrsstartstopstepaccessed_mappingr7)rtrWr8 __getitem__s     zSliceImageResult.__getitem__cCs t|jSrg)r|rbrhr7r7r8__len__s zSliceImageResult.__len__rg)r]r=r^r_)rcrR)r"rm)r"r=) rYrZr[rXrfpropertyrirkrnrorqrrr7r7r7r8r\s        r\皙?FrTstr | Image.Imagelist[CocoAnnotation] | Noneoutput_file_namer_ output_dirrMout_extverboseexif_fixc " sH| rtjnddd fdd } |d urt|jd d d t|d }dt|j|j\}}|dkr:|dksCtd|jdt|||||||d}d}t ||g|d}t |}|D]}|d7}|d}|d}|d}|d}|||||f}d t t|}| r| }n!t|drtt|dj}|tvrd}n |tvrt|jj}nd}|d||}|d|d}|d|d}t|||d}|d urt||| D]}||qt|||d|dgd} || q_|r|rtjjtd}!|! | |j|gt||jdt|dt|dt||S)!aSlice a large image into smaller windows. If output_file_name and output_dir is given, export sliced images. Args: image (str or PIL.Image): File path of image or Pillow Image to be sliced. coco_annotation_list (List[CocoAnnotation], optional): List of CocoAnnotation objects. output_file_name (str, optional): Root name of output files (coordinates will be appended to this) output_dir (str, optional): Output directory slice_height (int, optional): Height of each slice. Default None. slice_width (int, optional): Width of each slice. Default None. overlap_height_ratio (float, optional): Fractional overlap in height of each slice (e.g. an overlap of 0.2 for a slice of size 100 yields an overlap of 20 pixels). Default 0.2. overlap_width_ratio (float, optional): Fractional overlap in width of each slice (e.g. an overlap of 0.2 for a slice of size 100 yields an overlap of 20 pixels). Default 0.2. auto_slice_resolution (bool, optional): if not set slice parameters such as slice_height and slice_width, it enables automatically calculate these params from image resolution and orientation. min_area_ratio (float, optional): If the cropped annotation area to original annotation ratio is smaller than this value, the annotation is filtered out. Default 0.1. out_ext (str, optional): Extension of saved images. Default is the original suffix for lossless image formats and png for lossy formats ('.jpg','.jpeg'). verbose (bool, optional): Switch to print relevant values to screen. Default 'False'. exif_fix (bool): Whether to apply an EXIF fix to the image. Returns: sliced_image_result: SliceImageResult: sliced_image_list: list of SlicedImage image_dir: str Directory of the sliced image exports. original_image_size: list of int Size of the unsliced original image in [height, width] c_sdSrgr7)akr7r7r85szslice_image..rT np.ndarrayrstrslice_file_namecs>t|d}tt||}|||d|dS)Nrzsliced image path: )rrrsaveclose)rTrr image_pilslice_file_pathr verboselogr7r8_export_single_slice7s  z)slice_image.._export_single_sliceNT)parentsexist_okrz image.shape: rzinvalid image size: z for 'slice_image'.)rrrrrrr!)r]r^rBr@rA_rrz.png)rpr%r&rS) max_workersz Num slices: z slice_height: z slice_width: )rTrrrrr) r informkdirrrsize RuntimeErrorr9r\rwasarrayjoinrhasattrgetattrsuffixrrrrrrQadd_annotationrRrf concurrentfuturesThreadPoolExecutor MAX_WORKERSrkr|rq)"rTrHrrrrrr!rrMrrrrrrrr,n_imssliced_image_result image_pil_arrr<tlxtlybrxbryimage_pil_sliceslice_suffixesrrrUrPrc conc_execr7rr8 slice_images|3           "rcoco_annotation_file_pathrr^ output_coco_annotation_file_nameignore_negative_sampleslist[dict | str]c Cst|} t| }g}tt|jD]?\}}tj||j }z"t ||j t |j j d||||||| | | | d }||jWqtyQtd|Yqwt|| d|d} d}|rn|rnt ||d}t| || |fS)aSlice large images given in a directory, into smaller windows. If output_dir is given, export sliced images and coco file. Args: coco_annotation_file_path (str): Location of the coco annotation file image_dir (str): Base directory for the images output_coco_annotation_file_name (str): File name of the exported coco dataset json. output_dir (str, optional): Output directory ignore_negative_samples (bool, optional): If True, images without annotations are ignored. Defaults to False. slice_height (int, optional): Height of each slice. Default 512. slice_width (int, optional): Width of each slice. Default 512. overlap_height_ratio (float, optional): Fractional overlap in height of each slice (e.g. an overlap of 0.2 for a slice of size 100 yields an overlap of 20 pixels). Default 0.2. overlap_width_ratio (float, optional): Fractional overlap in width of each slice (e.g. an overlap of 0.2 for a slice of size 100 yields an overlap of 20 pixels). Default 0.2. min_area_ratio (float): If the cropped annotation area to original annotation ratio is smaller than this value, the annotation is filtered out. Default 0.1. out_ext (str, optional): Extension of saved images. Default is the original suffix. verbose (bool, optional): Switch to print relevant values to screen. exif_fix (bool, optional): Whether to apply an EXIF fix to the image. Returns: coco_dict: dict COCO dict for sliced images and annotations save_path: str Path to the saved coco file r) rTrHrrrrrr!rMrrrz/Invalid annotation found, skipping this image: categories)rz _coco.json)rr from_coco_dict_or_path enumeraterrkospathrrprrrstemextendrnrr warningrr)rr^rrrrrrr!rMrrr coco_dictcocosliced_coco_imagesidxrU image_pathslice_image_result save_pathr7r7r8 slice_cocos@1    rrB orientation+Literal['vertical', 'horizontal', 'square']slideratiofloatcCs~|dkr||d||f\}}}}n(|dkr"|d|||f\}}}}n|dkr1||||f\}}}}ntd|d||||fS)z According to image resolution calculation overlap params Args: orientation: image capture angle slide: sliding window ratio: buffer value Returns: overlap params verticalr@ horizontalsquarezInvalid orientation: z7. Must be one of 'vertical', 'horizontal', or 'square'.)r()rrr slice_row slice_colrr!r7r7r8calc_ratio_and_slices  r resolutioncCs4d}td||kr|d7}td||ks |dS)z According to image resolution calculate power(2,n) and return the closest smaller `n`. Args: resolution: the width and height of the image multiplied. such as 1024x720 = 737280 Returns: rr@rB)rwpower)rexpor7r7r8calc_resolution_factor s  rr&r%cCs||krdS||kr dSdS)z_ Args: width: height: Returns: image capture orientation rrrr7r&r%r7r7r8calc_aspect_ratio_orientations  rtuple[int, int, int, int]c Cs|dkrt|ddd\}}}}n(|dkr t|ddd\}}}}n|dkr0t|d dd\}}}}nd}d}d}d}||}||} t| |} t||} | | | |fS) a This function calculate according to image resolution slice and overlap params. Args: resolution: str height: int width: int orientation: str Returns: x_overlap, y_overlap, slice_width, slice_height mediumrBg?)rrhighr@g? ultra-high)rr) rr%r&r split_row split_colrr!rrr0r/r7r7r8calc_slice_and_overlap_params-s*   rrescCs0t||d}t||||d\}}}}||||fS)z Args: res: resolution of image such as low, medium height: width: Returns: trigger slicing params function and return overlap params r)rr%r&r)rr)rr%r&rr0r/rrr7r7r8get_resolution_selectorYs  rcCs||}t|}|dkrtd||dSd|krdkr&nntd||dSd|kr0dkr9nntd||dStd||dS) aW According to Image HxW calculate overlap sliding window and buffer params factor is the power value of 2 closest to the image resolution. factor <= 18: low resolution image such as 300x300, 640x640 18 < factor <= 21: medium resolution image such as 1024x1024, 1336x960 21 < factor <= 24: high resolution image such as 2048x2048, 2048x4096, 4096x4096 factor > 24: ultra-high resolution image such as 6380x6380, 4096x8192 Args: height: width: Returns: slicing overlap params x_overlap, y_overlap, slice_width, slice_height lowr$rrr)rr)r%r&rfactorr7r7r8r'lsr'offset Sequence[int]cCsg}t|jdkr d}nd}|D]}|st|tjr|}t||d}|}|| qt|tjr=tj |ddS|rD| |S|S)aMShift bboxes w.r.t offset. Suppo Args: bboxes (Tensor, np.ndarray, list): The bboxes need to be translated. Its shape can be (n, 4), which means (x, y, x, y). offset (Sequence[int]): The translation offsets with shape of (2, ). Returns: Tensor, np.ndarray, list: Shifted bboxes. torchTF) shift_amountraxis) rrZrdrwrxryr get_shifted_boxr+to_xyxystack new_tensor)bboxesrshifted_bboxesbboxes_is_torch_tensorr?r7r7r8 shift_bboxess    rmasksr full_shapecCsJ|dur|Sg}|D]}t|||d}|}||jq tj|ddS)aAShift masks to the original image. Args: masks (np.ndarray): masks that need to be shifted. offset (Sequence[int]): The offset to translate with shape of (2, ). full_shape (Sequence[int]): A (height, width) tuple of the huge image's shape. Returns: np.ndarray: Shifted masks. N) segmentationrrrr)r get_shifted_maskr+ bool_maskrwr)rrr shifted_masksmaskr7r7r8 shift_maskss r)NNTrr)rrrrrrrrrrrr r!r r"r#)r:r;r<r=r"r>)rHrIr<r=r"rI) NNNNNrrTrNFT)rTrrHrrr_rr_rrrrrr r!r rrrMr rr_rrrr>r"r\) NFrrrrrNFT)rrr^rrrrr_rrrrrrrr r!r rMr rr_rrrr>r"r)rBr)rrrrrr)rrr"r)r&rr%rr"r) rrr%rr&rrrr"r)rrr%rr&rr"r)r%rr&rr"r)rr)rrrrrrr"r)3 __future__rconcurrent.futuresrrcollections.abcrpathlibrtypingrnumpyrwPILrshapely.errorsrrsahi.annotationr r sahi.loggerr sahi.utils.cocor r rr sahi.utils.cvrrrsahi.utils.filerrrr9rGrQrRr\rrrrrrrr'rrr7r7r7r8sv         C e  ]    ,  "