From a35f8ccfafead166d7b62187ad77f55a0f3d61b7 Mon Sep 17 00:00:00 2001 From: Yagmur Gizem Cinar Date: Thu, 15 Feb 2024 15:50:30 -0500 Subject: [PATCH] refactor: Updating image segmentation module inference Signed-off-by: Yagmur Gizem Cinar --- .../data_model/image_segmentation.py | 4 ++- .../data_model/object_detection.py | 8 ++--- .../modules/segmentation/image_processing.py | 29 +++++++++++-------- .../modules/segmentation/models.py | 10 ++----- pyproject.toml | 8 ++--- .../segmentation/test_segmentation_model.py | 3 +- 6 files changed, 31 insertions(+), 31 deletions(-) diff --git a/caikit_computer_vision/data_model/image_segmentation.py b/caikit_computer_vision/data_model/image_segmentation.py index 6680f57..72182f4 100644 --- a/caikit_computer_vision/data_model/image_segmentation.py +++ b/caikit_computer_vision/data_model/image_segmentation.py @@ -41,7 +41,9 @@ class ObjectSegment(DataObjectBase): # This mask should be image mode (L), 8 bit grayscale image treated # as a binary image, where 0 is background, and 255 is part of the # object to align with HF task definitions. - mask: Annotated[Union[caikit_dm.Image, List[float]], FieldNumber(5)] + # mask or polygon -- one of them is returned + polygon: Annotated[List[float], FieldNumber(5)] + mask: Annotated[caikit_dm.Image, FieldNumber(6)] @dataobject(package="caikit_data_model.caikit_computer_vision") diff --git a/caikit_computer_vision/data_model/object_detection.py b/caikit_computer_vision/data_model/object_detection.py index 4a76653..0a7c72e 100644 --- a/caikit_computer_vision/data_model/object_detection.py +++ b/caikit_computer_vision/data_model/object_detection.py @@ -42,10 +42,10 @@ class Point2d(DataObjectBase): @dataobject(package="caikit_data_model.caikit_computer_vision") class BoundingBox(DataObjectBase): - xmin: Annotated[int, FieldNumber(1)] - xmax: Annotated[int, FieldNumber(2)] - ymin: Annotated[int, FieldNumber(3)] - ymax: Annotated[int, FieldNumber(4)] + xmin: Annotated[float, FieldNumber(1)] + xmax: Annotated[float, FieldNumber(2)] + ymin: Annotated[float, FieldNumber(3)] + ymax: Annotated[float, FieldNumber(4)] @dataobject(package="caikit_data_model.caikit_computer_vision") diff --git a/caikit_computer_vision/modules/segmentation/image_processing.py b/caikit_computer_vision/modules/segmentation/image_processing.py index bd57bc1..8cd9073 100644 --- a/caikit_computer_vision/modules/segmentation/image_processing.py +++ b/caikit_computer_vision/modules/segmentation/image_processing.py @@ -50,7 +50,8 @@ def maskrcnn_postprocess( results (dict): COCO json from an inference task uuid (str): Unique identifier of image image_id (int, optional): Image ID to store in the task. Defaults to 1. - + threshold (float): Threshold in range (0,1] to be used for filtering bounding box predictions. + Default 0.2 Returns: dict: Formatted COCO JSON in OCL format """ @@ -62,8 +63,12 @@ def maskrcnn_postprocess( try: result_instances = detector_postprocess(instances, image_size[0], image_size[1]) except IndexError as e: - log.error(f"Length of MASKS {len(results['MASKS'])}") - log.error("No instances found, setting empty annotations", exc_info=e) + error("", f"Length of MASKS {len(results['MASKS'])}") + error( + "", + "No instances found, setting empty annotations", + exc_info=e, + ) result_instances = [] coco_instances = instances_to_coco_json(result_instances, image_id) coco_json = _assemble_images_coco( @@ -167,11 +172,13 @@ def detectron2_to_coco( if isinstance(inf_coco, str): inf_coco = COCO(inf_coco) else: - assert isinstance(inf_coco, COCO), "inf_coco must be an instance of class COCO" + error.type_check("", COCO, inf_coco=inf_coco) - assert ( - "categories" and "images" in inf_coco.dataset - ), "inf_coco should contain both 'categories' and 'images' keys" + error.value_check( + "", + ("categories" and "images" in inf_coco.dataset), + "inf_coco should contain both 'categories' and 'images' keys", + ) # give ids to annotations # correct the placement of the score key as in coco @@ -230,10 +237,10 @@ def convert_seg_poly(inf_coco: Type[COCO]) -> Type[COCO]: """converts segmentation from detectron2 format to polygons. Args: - inf_coco : coco object with annotations containing detectron2 formatted segmentations + inf_coco : coco object with annotations dictionary containing detectron2 formatted segmentations Returns: - inf_coco : coco dict with annotations containing coco formatted polygons + inf_coco : coco object with annotations dictionary containing coco formatted polygons """ def mask_to_polygons(mask): @@ -245,14 +252,12 @@ def mask_to_polygons(mask): mask = np.ascontiguousarray( mask ) # some versions of cv2 does not support incontiguous arr - res = cv2.findContours( + res, hierarchy = cv2.findContours( mask.astype("uint8"), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE ) - hierarchy = res[-1] if hierarchy is None: # empty mask return [], False has_holes = (hierarchy.reshape(-1, 4)[:, 3] >= 0).sum() > 0 - res = res[-2] res = [x.flatten() for x in res] # These coordinates from OpenCV are integers in range [0, W-1 or H-1]. # We add 0.5 to turn them into real-value coordinate space. A better solution diff --git a/caikit_computer_vision/modules/segmentation/models.py b/caikit_computer_vision/modules/segmentation/models.py index f1e2049..02444c3 100644 --- a/caikit_computer_vision/modules/segmentation/models.py +++ b/caikit_computer_vision/modules/segmentation/models.py @@ -84,24 +84,18 @@ def load(cls, model_file: Union[str, "ModuleConfig"]) -> "ViTSegmenter": def save( self, model_path: str, - segmentation_dirname: str = "image_segmentation", ): """Save the in-memory model to the given path. Args: model_path: str Path that we want to export the segmentation model to. - segmentation_dirname: str - Subdirectory to which we want to save the segmentation model. - Default: image_segmentation """ saver = ModuleSaver( self, model_path=model_path, ) - segmentation_rel_path, segmentation_abs_path = saver.add_dir( - segmentation_dirname - ) + segmentation_rel_path, segmentation_abs_path = saver.add_dir(model_path) with saver: saver.update_config( { @@ -181,7 +175,7 @@ def run( score=results[idx]["attributes"]["score"], category_id=results[idx]["category_id"], bbox=BoundingBox(*results[idx]["bbox"]), - mask=results[idx]["segmentation"], + polygon=results[idx]["segmentation"], area=results[idx]["area"], ) for idx in range(num_objects) diff --git a/pyproject.toml b/pyproject.toml index fc2d755..58668b1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,10 +18,10 @@ dependencies = [ "transformers>=4.27.1,<5", "torch>2.0,<3", "timm>=0.9.5,<1", - "opencv-python", - "pycocotools", - "detectron2 @git+https://github.com/facebookresearch/detectron2.git", - "shapely" + "opencv-python>=4.9.0.80,<5", + "pycocotools>=2.0.7,<3", + "detectron2 @git+https://github.com/facebookresearch/detectron2.git@e70b9229d77aa39d85f8fa5266e6ea658e92eed3", + "shapely>=2.0.2,<3", ] [project.urls] diff --git a/tests/modules/segmentation/test_segmentation_model.py b/tests/modules/segmentation/test_segmentation_model.py index 06b73cc..5ffc737 100644 --- a/tests/modules/segmentation/test_segmentation_model.py +++ b/tests/modules/segmentation/test_segmentation_model.py @@ -37,8 +37,7 @@ def test_save_model_and_bootstrap(): """Ensure that we can save a model and reload it.""" model = ViTSegmenter.bootstrap(SEGMENTATION_MODEL_CKPT) with tempfile.TemporaryDirectory() as model_dir: - model.save(model_dir, "segmenter") - del model + model.save(model_dir) new_model = ViTSegmenter.load(model_dir) preds = new_model.run(np.ones((800, 800, 3), dtype=np.uint8)) assert isinstance(preds, ImageSegmentationResult)