diff --git a/sahi/slicing.py b/sahi/slicing.py index ebeeba4b7..a8c1d7d91 100644 --- a/sahi/slicing.py +++ b/sahi/slicing.py @@ -120,7 +120,9 @@ def annotation_inside_slice(annotation: Dict, slice_bbox: List[int]) -> bool: return True -def process_coco_annotations(coco_annotation_list: List[CocoAnnotation], slice_bbox: List[int], min_area_ratio) -> bool: +def process_coco_annotations( + coco_annotation_list: List[CocoAnnotation], slice_bbox: List[int], min_area_ratio +) -> List[CocoAnnotation]: """Slices and filters given list of CocoAnnotation objects with given 'slice_bbox' and 'min_area_ratio'. @@ -161,19 +163,19 @@ def __init__(self, image, coco_image, starting_pixel): class SliceImageResult: - def __init__(self, original_image_size=None, image_dir: str = None): + def __init__(self, original_image_size: List[int], image_dir: Optional[str] = None): """ - 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] """ - self._sliced_image_list: List[SlicedImage] = [] self.original_image_height = original_image_size[0] self.original_image_width = original_image_size[1] self.image_dir = image_dir + self._sliced_image_list: List[SlicedImage] = [] + def add_sliced_image(self, sliced_image: SlicedImage): if not isinstance(sliced_image, SlicedImage): raise TypeError("sliced_image must be a SlicedImage instance") @@ -261,13 +263,13 @@ def __len__(self): def slice_image( image: Union[str, Image.Image], - coco_annotation_list: Optional[CocoAnnotation] = None, + coco_annotation_list: Optional[List[CocoAnnotation]] = None, output_file_name: Optional[str] = None, output_dir: Optional[str] = None, - slice_height: int = None, - slice_width: int = None, - overlap_height_ratio: float = None, - overlap_width_ratio: float = None, + slice_height: Optional[int] = None, + slice_width: Optional[int] = None, + overlap_height_ratio: float = 0.2, + overlap_width_ratio: float = 0.2, auto_slice_resolution: bool = True, min_area_ratio: float = 0.1, out_ext: Optional[str] = None, @@ -278,12 +280,12 @@ def slice_image( Args: image (str or PIL.Image): File path of image or Pillow Image to be sliced. - coco_annotation_list (CocoAnnotation): List of CocoAnnotation objects. + 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): Height of each slice. Default 512. - slice_width (int): Width of each slice. Default 512. + slice_height (int, optional): Height of each slice. Default None. + slice_width (int, optional): Width of each slice. Default None. overlap_height_ratio (float): 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. @@ -342,7 +344,6 @@ def _export_single_slice(image: np.ndarray, output_dir: str, slice_file_name: st overlap_width_ratio=overlap_width_ratio, ) - t0 = time.time() n_ims = 0 # init images and annotations lists @@ -360,23 +361,18 @@ def _export_single_slice(image: np.ndarray, output_dir: str, slice_file_name: st bry = slice_bbox[3] image_pil_slice = image_pil_arr[tly:bry, tlx:brx] - # process annotations if coco_annotations is given - if coco_annotation_list is not None: - sliced_coco_annotation_list = process_coco_annotations(coco_annotation_list, slice_bbox, min_area_ratio) - # set image file suffixes slice_suffixes = "_".join(map(str, slice_bbox)) if out_ext: suffix = out_ext - else: - try: - suffix = Path(image_pil.filename).suffix - if suffix in IMAGE_EXTENSIONS_LOSSY: - suffix = ".png" - elif suffix in IMAGE_EXTENSIONS_LOSSLESS: - suffix = Path(image_pil.filename).suffix - except AttributeError: + elif hasattr(image_pil, "filename"): + suffix = Path(getattr(image_pil, "filename")).suffix + if suffix in IMAGE_EXTENSIONS_LOSSY: suffix = ".png" + elif suffix in IMAGE_EXTENSIONS_LOSSLESS: + suffix = Path(image_pil.filename).suffix + else: + suffix = ".png" # set image file name and path slice_file_name = f"{output_file_name}_{slice_suffixes}{suffix}" @@ -387,9 +383,9 @@ def _export_single_slice(image: np.ndarray, output_dir: str, slice_file_name: st coco_image = CocoImage(file_name=slice_file_name, height=slice_height, width=slice_width) # append coco annotations (if present) to coco image - if coco_annotation_list: - for coco_annotation in sliced_coco_annotation_list: - coco_image.add_annotation(coco_annotation) + if coco_annotation_list is not None: + for sliced_coco_annotation in process_coco_annotations(coco_annotation_list, slice_bbox, min_area_ratio): + coco_image.add_annotation(sliced_coco_annotation) # create sliced image and append to sliced_image_result sliced_image = SlicedImage(