Add visual prompting zero-shot learning (export, IR inference)

CHANGELOG.md

@@ -6,7 +6,7 @@ All notable changes to this project will be documented in this file.
 
 ### New features
 
-- Add zero-shot visual prompting (https://github.com/openvinotoolkit/training_extensions/pull/2616)
+- Add zero-shot visual prompting (<https://github.com/openvinotoolkit/training_extensions/pull/2616>, <https://github.com/openvinotoolkit/training_extensions/pull/2706>)
 
 ### Enhancements
 

src/otx/algorithms/visual_prompting/adapters/openvino/model_wrappers/__init__.py

@@ -14,4 +14,4 @@
 # See the License for the specific language governing permissions
 # and limitations under the License.
 
-from .openvino_models import Decoder, ImageEncoder  # noqa: F401
+from .openvino_models import Decoder, ImageEncoder, PromptGetter  # noqa: F401

src/otx/algorithms/visual_prompting/adapters/openvino/model_wrappers/openvino_models.py

@@ -59,6 +59,20 @@ def preprocess(
         return dict_inputs, meta
 
 
+class PromptGetter(ImageModel):
+    """PromptGetter class for zero-shot visual prompting of openvino model wrapper."""
+
+    __model__ = "prompt_getter"
+
+    @classmethod
+    def parameters(cls) -> Dict[str, Any]:  # noqa: D102
+        parameters = super().parameters()
+        parameters.update({"image_size": NumericalValue(value_type=int, default_value=1024, min=0, max=2048)})
+        parameters.update({"sim_threshold": NumericalValue(value_type=float, default_value=0.5, min=0, max=1)})
+        parameters.update({"num_bg_points": NumericalValue(value_type=int, default_value=1, min=0, max=1024)})
+        return parameters
+
+
 class Decoder(SegmentationModel):
     """Decoder class for visual prompting of openvino model wrapper."""
 
@@ -76,6 +90,7 @@ def __init__(
     def parameters(cls):  # noqa: D102
         parameters = super().parameters()
         parameters.update({"image_size": NumericalValue(value_type=int, default_value=1024, min=0, max=2048)})
+        parameters.update({"mask_threshold": NumericalValue(value_type=float, default_value=0.0, min=0, max=1)})
         return parameters
 
     def _get_outputs(self):
@@ -174,7 +189,7 @@ def resize_and_crop(self, soft_prediction: np.ndarray, original_size: np.ndarray
         )
 
         prepadded_size = self.get_padded_size(original_size, self.image_size).astype(np.int64)
-        resized_cropped_soft_prediction = resized_soft_prediction[..., : prepadded_size[0], : prepadded_size[1]]
+        resized_cropped_soft_prediction = resized_soft_prediction[: prepadded_size[0], : prepadded_size[1], ...]
 
         original_size = original_size.astype(np.int64)
         h, w = original_size

src/otx/algorithms/visual_prompting/adapters/pytorch_lightning/datasets/dataset.py

@@ -207,7 +207,7 @@ def get_prompts(dataset_item: DatasetItemEntity, dataset_labels: List[LabelEntit
 
         bboxes = np.array(bboxes)
         return dict(
-            original_size=(height, width),
+            original_size=np.array((height, width), dtype=np.int64),
             gt_masks=gt_masks,
             bboxes=bboxes,
             points=points,  # TODO (sungchul): update point information
@@ -247,6 +247,20 @@ def __getitem__(self, index: int) -> Dict[str, Union[int, List, Tensor]]:
 class OTXZeroShotVisualPromptingDataset(OTXVisualPromptingDataset):
     """Visual Prompting for Zero-shot learning Dataset Adaptor."""
 
+    def __init__(
+        self,
+        dataset: DatasetEntity,
+        image_size: int,
+        mean: List[float],
+        std: List[float],
+        generate_point: bool = False,
+        generate_bbox: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__(dataset, image_size, mean, std, offset_bbox=0)
+        self.generate_point = generate_point
+        self.generate_bbox = generate_bbox
+
     def __getitem__(self, index: int) -> Dict[str, Union[int, List, Tensor]]:
         """Get dataset item.
 
@@ -288,7 +302,7 @@ def __init__(
         self.config = config
         self.dataset = dataset
         self.train_type = train_type
-        # self.kwargs = {}
+        self.kwargs = {}
         if self.train_type == TrainType.Zeroshot:
             # check zero-shot configs
             if self.config.get("train_batch_size", 1) != 1:
@@ -300,12 +314,12 @@ def __init__(
                 )
                 self.config["train_batch_size"] = 1
 
-            # self.kwargs.update(
-            #     {
-            #         "generate_point": self.config.get("generate_point", False),
-            #         "generate_bbox": self.config.get("generate_bbox", False),
-            #     }
-            # )
+            self.kwargs.update(
+                {
+                    "generate_point": self.config.get("generate_point", False),
+                    "generate_bbox": self.config.get("generate_bbox", False),
+                }
+            )
 
         self.train_otx_dataset: DatasetEntity
         self.val_otx_dataset: DatasetEntity
@@ -331,7 +345,7 @@ def setup(self, stage: Optional[str] = None) -> None:
                 mean=mean,
                 std=std,
                 offset_bbox=self.config.offset_bbox,
-                # **self.kwargs,
+                **self.kwargs,
             )
 
             # self.val_dataset = None
@@ -347,11 +361,7 @@ def setup(self, stage: Optional[str] = None) -> None:
 
         if stage == "predict":
             self.predict_dataset = self.DATASETS[self.train_type](
-                dataset=self.dataset,
-                image_size=image_size,
-                mean=mean,
-                std=std,
-                # **self.kwargs
+                dataset=self.dataset, image_size=image_size, mean=mean, std=std, **self.kwargs
             )
 
     def summary(self):

src/otx/algorithms/visual_prompting/adapters/pytorch_lightning/datasets/pipelines/sam_transforms.py

@@ -4,7 +4,6 @@
 # All rights reserved.
 #
 
-from copy import deepcopy
 from typing import Any, Dict, List, Tuple, Union
 
 import numpy as np
@@ -76,9 +75,9 @@ def apply_coords(
         old_h, old_w = original_size
         new_h, new_w = cls.get_preprocess_shape(original_size[0], original_size[1], target_length)
         if isinstance(coords, np.ndarray):
-            coords = deepcopy(coords).astype(np.float32)
+            coords = coords.astype(float)
         else:
-            coords = deepcopy(coords).to(torch.float32)
+            coords = coords.to(torch.float)
         coords[..., 0] = coords[..., 0] * (new_w / old_w)
         coords[..., 1] = coords[..., 1] * (new_h / old_h)
         return coords

src/otx/algorithms/visual_prompting/adapters/pytorch_lightning/datasets/pipelines/transforms.py

@@ -23,25 +23,26 @@ def collate_fn(batch: List[Any]) -> Dict:
         Dict: Collated batch data.
     """
 
-    def _convert_empty_to_none(x: str) -> List:
+    def _convert_empty_to_none(x: str, dtype: torch.dtype = torch.float32) -> List:
         """Convert empty list to None.
 
         Args:
             x (str): Key of batch data.
+            dtype (torch.dtype): Dtype to be applied to tensors.
 
         Returns:
             List: List of batch data.
         """
         func = torch.stack if x == "gt_masks" else torch.tensor
-        items = [func(item[x]) for item in batch if item[x] is not None]
+        items = [func(item[x]).to(dtype) for item in batch if item[x] is not None]
         return None if len(items) == 0 else items
 
     index = [item["index"] for item in batch]
     images = torch.stack([item["images"] for item in batch])
     bboxes = _convert_empty_to_none("bboxes")
     points = None  # TBD
-    gt_masks = _convert_empty_to_none("gt_masks")
-    original_size = [item["original_size"] for item in batch]
+    gt_masks = _convert_empty_to_none("gt_masks", torch.int32)
+    original_size = _convert_empty_to_none("original_size")
     padding = [item["padding"] for item in batch]
     path = [item["path"] for item in batch]
     labels = [item["labels"] for item in batch]

src/otx/algorithms/visual_prompting/adapters/pytorch_lightning/models/visual_prompters/segment_anything.py

@@ -10,7 +10,7 @@
 
 import re
 from collections import OrderedDict
-from typing import Dict, List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from omegaconf import DictConfig
@@ -334,47 +334,36 @@ def select_masks(self, masks: Tensor, iou_preds: Tensor, num_points: int) -> Tup
 
         return masks, iou_preds
 
-    def mask_postprocessing(self, masks: Tensor, orig_size: Tensor) -> Tensor:
+    @staticmethod
+    def mask_postprocessing(masks: Tensor, input_size: int, orig_size: Tensor) -> Tensor:
         """Postprocesses the predicted masks.
 
         Args:
             masks (Tensor): A batch of predicted masks with shape Bx1xHxW.
+            input_size (int): The size of the image input to the model, in (H, W) format.
+                Used to remove padding.
             orig_size (Tensor): The original image size with shape Bx2.
 
         Returns:
             masks (Tensor): The postprocessed masks with shape Bx1xHxW.
         """
-        masks = F.interpolate(
-            masks,
-            size=(self.config.model.image_size, self.config.model.image_size),
-            mode="bilinear",
-            align_corners=False,
-        )
 
-        prepadded_size = self.resize_longest_image_size(orig_size, self.config.model.image_size).to(torch.int64)
+        def resize_longest_image_size(input_image_size: Tensor, longest_side: int) -> Tensor:
+            scale = longest_side / torch.max(input_image_size)
+            transformed_size = scale * input_image_size
+            transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64)
+            return transformed_size
+
+        masks = F.interpolate(masks, size=(input_size, input_size), mode="bilinear", align_corners=False)
+
+        prepadded_size = resize_longest_image_size(orig_size, input_size)
         masks = masks[..., : prepadded_size[0], : prepadded_size[1]]  # type: ignore
 
         orig_size = orig_size.to(torch.int64)
         h, w = orig_size[0], orig_size[1]
         masks = F.interpolate(masks, size=(h, w), mode="bilinear", align_corners=False)
         return masks
 
-    def resize_longest_image_size(self, input_image_size: Tensor, longest_side: int) -> Tensor:
-        """Resizes the longest side of the image to the given size.
-
-        Args:
-            input_image_size (Tensor): The original image size with shape Bx2.
-            longest_side (int): The size of the longest side.
-
-        Returns:
-            transformed_size (Tensor): The transformed image size with shape Bx2.
-        """
-        input_image_size = input_image_size.to(torch.float32)
-        scale = longest_side / torch.max(input_image_size)
-        transformed_size = scale * input_image_size
-        transformed_size = torch.floor(transformed_size + 0.5).to(torch.int64)
-        return transformed_size
-
     ######################################################
     #     forward for training/validation/prediction     #
     ######################################################
@@ -556,26 +545,26 @@ def predict_step(self, batch, batch_idx) -> Dict[str, Tensor]:
     def postprocess_masks(
         masks: Tensor,
         input_size: Tuple[int, int],
-        padding: Tuple[int, ...],
-        original_size: Tuple[int, int],
+        padding: Union[Tuple[int, ...], Tensor],
+        original_size: Union[Tuple[int, int], Tensor],
     ) -> Tensor:
         """Remove padding and upscale masks to the original image size.
 
         Args:
             masks (Tensor): Predicted masks from the mask_decoder with (N, 1, H/downsized_ratio, W/downsized_ratio).
             input_size (tuple(int, int)): The size of the image input to the model, in (H, W) format.
                 Used to remove padding.
-            padding (tuple(int, int, int, int), optional): The padding applied to the image before input to the model,
+            padding (tuple(int, int, int, int), Tensor): The padding applied to the image before input to the model,
                 in (left, top, right, bottom) format.
-            original_size (tuple(int, int)): The original size of the image before resizing for input to the model,
-                in (H, W) format.
+            original_size (tuple(int, int), Tensor): The original size of the image before resizing
+                for input to the model, in (H, W) format.
 
         Returns:
           (Tensor): Postprocessed masks in NxHxW format, where (H, W) is given by original_size.
         """
         masks = F.interpolate(masks, input_size, mode="bilinear", align_corners=False)
         masks = masks[..., : input_size[0] - padding[3], : input_size[1] - padding[2]]
-        masks = F.interpolate(masks, original_size, mode="bilinear", align_corners=False)
+        masks = F.interpolate(masks, [int(o) for o in original_size], mode="bilinear", align_corners=False)
         return masks.squeeze(1)
 
     def configure_optimizers(self) -> optim: