onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/models/sam2/sam2_demo.py

@@ -0,0 +1,321 @@
+# -------------------------------------------------------------------------
+# Copyright (R) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import os
+
+import matplotlib.image as mpimg
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+from matplotlib.patches import Rectangle
+from PIL import Image
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+from sam2_image_onnx_predictor import SAM2ImageOnnxPredictor
+from sam2_utils import load_sam2_model
+
+import onnxruntime
+
+
+def show_mask(mask, ax, random_color=False, borders=True):
+    if random_color:
+        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
+    else:
+        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
+    h, w = mask.shape[-2:]
+    mask = mask.astype(np.uint8)
+    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
+    if borders:
+        import cv2  # noqa: PLC0415
+
+        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+        # Try to smooth contours
+        contours = [cv2.approxPolyDP(contour, epsilon=0.01, closed=True) for contour in contours]
+        mask_image = cv2.drawContours(mask_image, contours, -1, (1, 1, 1, 0.5), thickness=2)
+    ax.imshow(mask_image)
+
+
+def show_points(coords, labels, ax, marker_size=375):
+    pos_points = coords[labels == 1]
+    neg_points = coords[labels == 0]
+    ax.scatter(
+        pos_points[:, 0], pos_points[:, 1], color="green", marker="*", s=marker_size, edgecolor="white", linewidth=1.25
+    )
+    ax.scatter(
+        neg_points[:, 0], neg_points[:, 1], color="red", marker="*", s=marker_size, edgecolor="white", linewidth=1.25
+    )
+
+
+def show_box(box, ax):
+    x0, y0 = box[0], box[1]
+    w, h = box[2] - box[0], box[3] - box[1]
+    ax.add_patch(Rectangle((x0, y0), w, h, edgecolor="green", facecolor=(0, 0, 0, 0), lw=2))
+
+
+def show_masks(
+    image,
+    masks,
+    scores,
+    point_coords=None,
+    box_coords=None,
+    input_labels=None,
+    borders=True,
+    output_image_file_prefix=None,
+    image_files=None,
+):
+    for i, (mask, score) in enumerate(zip(masks, scores, strict=False)):
+        plt.figure(figsize=(10, 10))
+        plt.imshow(image)
+        show_mask(mask, plt.gca(), borders=borders)
+        if point_coords is not None:
+            assert input_labels is not None
+            show_points(point_coords, input_labels, plt.gca())
+
+        if box_coords is not None:
+            show_box(box_coords, plt.gca())
+
+        if len(scores) > 1:
+            plt.title(f"Mask {i + 1}, Score: {score:.3f}", fontsize=18)
+
+        plt.axis("off")
+        if output_image_file_prefix:
+            filename = f"{output_image_file_prefix}_{i}.png"
+            if os.path.exists(filename):
+                os.remove(filename)
+            plt.savefig(filename, format="png", bbox_inches="tight", pad_inches=0)
+            if isinstance(image_files, list):
+                image_files.append(filename)
+        plt.show(block=False)
+        plt.close()
+
+
+def get_predictor(
+    sam2_dir: str,
+    device: str | torch.device,
+    dtype: torch.dtype,
+    model_type="sam2_hiera_large",
+    engine="torch",
+    image_encoder_onnx_path: str = "",
+    image_decoder_onnx_path: str = "",
+    image_decoder_multi_onnx_path: str = "",
+    provider: str = "CUDAExecutionProvider",
+):
+    sam2_model = load_sam2_model(sam2_dir, model_type, device=device)
+    if engine == "torch":
+        predictor = SAM2ImagePredictor(sam2_model)
+    else:
+        predictor = SAM2ImageOnnxPredictor(
+            sam2_model,
+            image_encoder_onnx_path=image_encoder_onnx_path,
+            image_decoder_onnx_path=image_decoder_onnx_path,
+            image_decoder_multi_onnx_path=image_decoder_multi_onnx_path,
+            provider=provider,
+            device=device,
+            onnx_dtype=dtype,
+        )
+    return predictor
+
+
+def run_demo(
+    sam2_dir: str,
+    model_type: str = "sam2_hiera_large",
+    engine: str = "torch",
+    dtype: torch.dtype = torch.float32,
+    image_encoder_onnx_path: str = "",
+    image_decoder_onnx_path: str = "",
+    image_decoder_multi_onnx_path: str = "",
+    use_gpu: bool = True,
+    enable_batch: bool = False,
+):
+    if use_gpu:
+        assert torch.cuda.is_available()
+        assert "CUDAExecutionProvider" in onnxruntime.get_available_providers()
+        provider = "CUDAExecutionProvider"
+    else:
+        provider = "CPUExecutionProvider"
+
+    device = torch.device("cuda" if use_gpu else "cpu")
+
+    if use_gpu and engine == "torch" and torch.cuda.get_device_properties(0).major >= 8:
+        # Turn on tfloat32 for Ampere GPUs.
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+    np.random.seed(3)
+    image = Image.open("truck.jpg")
+    image = np.array(image.convert("RGB"))
+
+    predictor = get_predictor(
+        sam2_dir,
+        device,
+        dtype,
+        model_type,
+        engine,
+        image_encoder_onnx_path,
+        image_decoder_onnx_path,
+        image_decoder_multi_onnx_path,
+        provider=provider,
+    )
+
+    predictor.set_image(image)
+    prefix = f"sam2_demo_{engine}_"
+
+    #  The model returns masks, quality predictions for those masks,
+    #     and low resolution mask logits that can be passed to the next iteration of prediction.
+    #  With multimask_output=True (the default setting), SAM 2 outputs 3 masks, where
+    #     scores gives the model's own estimation of the quality of these masks.
+    #  For ambiguous prompts such as a single point, it is recommended to use multimask_output=True
+    #     even if only a single mask is desired;
+    input_point = np.array([[500, 375]])
+    input_label = np.array([1])
+    masks, scores, logits = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        multimask_output=True,
+    )
+
+    sorted_ind = np.argsort(scores)[::-1]
+    masks = masks[sorted_ind]
+    scores = scores[sorted_ind]
+    logits = logits[sorted_ind]
+
+    image_files = []
+    show_masks(
+        image,
+        masks,
+        scores,
+        point_coords=input_point,
+        input_labels=input_label,
+        borders=True,
+        output_image_file_prefix=prefix + "multimask",
+        image_files=image_files,
+    )
+
+    # Multiple points.
+    input_point = np.array([[500, 375], [1125, 625]])
+    input_label = np.array([1, 1])
+    mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask
+    masks, scores, _ = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        mask_input=mask_input[None, :, :],
+        multimask_output=False,
+    )
+    show_masks(
+        image,
+        masks,
+        scores,
+        point_coords=input_point,
+        input_labels=input_label,
+        output_image_file_prefix=prefix + "multi_points",
+        image_files=image_files,
+    )
+
+    # Specify a window and a background point.
+    input_point = np.array([[500, 375], [1125, 625]])
+    input_label = np.array([1, 0])
+    mask_input = logits[np.argmax(scores), :, :]  # Choose the model's best mask
+    masks, scores, _ = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        mask_input=mask_input[None, :, :],
+        multimask_output=False,
+    )
+    show_masks(
+        image,
+        masks,
+        scores,
+        point_coords=input_point,
+        input_labels=input_label,
+        output_image_file_prefix=prefix + "background_point",
+        image_files=image_files,
+    )
+
+    # Take a box as input
+    input_box = np.array([425, 600, 700, 875])
+    masks, scores, _ = predictor.predict(
+        point_coords=None,
+        point_labels=None,
+        box=input_box[None, :],
+        multimask_output=False,
+    )
+    show_masks(
+        image,
+        masks,
+        scores,
+        box_coords=input_box,
+        output_image_file_prefix=prefix + "box",
+        image_files=image_files,
+    )
+
+    # Combining points and boxes
+    input_box = np.array([425, 600, 700, 875])
+    input_point = np.array([[575, 750]])
+    input_label = np.array([0])
+
+    masks, scores, logits = predictor.predict(
+        point_coords=input_point,
+        point_labels=input_label,
+        box=input_box,
+        multimask_output=False,
+    )
+    show_masks(
+        image,
+        masks,
+        scores,
+        box_coords=input_box,
+        point_coords=input_point,
+        input_labels=input_label,
+        output_image_file_prefix=prefix + "box_and_point",
+        image_files=image_files,
+    )
+
+    # TODO: support batched prompt inputs
+    if enable_batch:
+        input_boxes = np.array(
+            [
+                [75, 275, 1725, 850],
+                [425, 600, 700, 875],
+                [1375, 550, 1650, 800],
+                [1240, 675, 1400, 750],
+            ]
+        )
+        masks, scores, _ = predictor.predict(
+            point_coords=None,
+            point_labels=None,
+            box=input_boxes,
+            multimask_output=False,
+        )
+        plt.figure(figsize=(10, 10))
+        plt.imshow(image)
+        for mask in masks:
+            show_mask(mask.squeeze(0), plt.gca(), random_color=True)
+        for box in input_boxes:
+            show_box(box, plt.gca())
+        plt.axis("off")
+        plt.show()
+        plt.savefig(prefix + "batch_prompt.png")
+        image_files.append(prefix + "batch_prompt.png")
+    return image_files
+
+
+def show_all_images(left_images, right_images, suffix=""):
+    # Show images in two rows since display screen is horizontal in most cases.
+    fig, axes = plt.subplots(nrows=2, ncols=len(left_images), figsize=(19.20, 10.80))
+    for i, (left_img_path, right_img_path) in enumerate(zip(left_images, right_images, strict=False)):
+        left_img = mpimg.imread(left_img_path)
+        right_img = mpimg.imread(right_img_path)
+
+        axes[0, i].imshow(left_img)
+        axes[0, i].set_title(left_img_path.replace("sam2_demo_", "").replace(".png", ""), fontsize=10)
+        axes[0, i].axis("off")
+        axes[0, i].set_aspect(left_img.shape[1] / left_img.shape[0])
+
+        axes[1, i].imshow(right_img)
+        axes[1, i].set_title(right_img_path.replace("sam2_demo_", "").replace(".png", ""), fontsize=10)
+        axes[1, i].axis("off")
+        axes[1, i].set_aspect(right_img.shape[1] / right_img.shape[0])
+
+    plt.tight_layout()
+    plt.savefig(f"sam2_demo{suffix}.png", format="png", bbox_inches="tight", dpi=1000)
+    plt.show()

onnxruntime/transformers/models/sam2/sam2_image_onnx_predictor.py

@@ -0,0 +1,279 @@
+# -------------------------------------------------------------------------
+# Copyright (R) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import logging
+
+import numpy as np
+import torch
+from PIL.Image import Image
+from sam2.modeling.sam2_base import SAM2Base
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+from sam2_utils import decoder_shape_dict, encoder_shape_dict
+
+from onnxruntime import InferenceSession
+from onnxruntime.transformers.io_binding_helper import CudaSession
+
+logger = logging.getLogger(__name__)
+
+
+def create_ort_session(
+    onnx_path: str,
+    session_options=None,
+    provider="CUDAExecutionProvider",
+    enable_cuda_graph=False,
+    use_tf32=True,
+) -> InferenceSession:
+    if provider == "CUDAExecutionProvider":
+        device_id = torch.cuda.current_device()
+        provider_options = CudaSession.get_cuda_provider_options(device_id, enable_cuda_graph)
+        provider_options["use_tf32"] = int(use_tf32)
+        providers = [(provider, provider_options), "CPUExecutionProvider"]
+    else:
+        providers = ["CPUExecutionProvider"]
+    logger.info("Using providers: %s", providers)
+    return InferenceSession(onnx_path, session_options, providers=providers)
+
+
+def create_session(
+    onnx_path: str,
+    session_options=None,
+    provider="CUDAExecutionProvider",
+    device: str | torch.device = "cuda",
+    enable_cuda_graph=False,
+) -> CudaSession:
+    ort_session = create_ort_session(
+        onnx_path, session_options, provider, enable_cuda_graph=enable_cuda_graph, use_tf32=True
+    )
+    cuda_session = CudaSession(ort_session, device=torch.device(device), enable_cuda_graph=enable_cuda_graph)
+    return cuda_session
+
+
+class SAM2ImageOnnxPredictor(SAM2ImagePredictor):
+    def __init__(
+        self,
+        sam_model: SAM2Base,
+        image_encoder_onnx_path: str = "",
+        image_decoder_onnx_path: str = "",
+        image_decoder_multi_onnx_path: str = "",
+        provider: str = "CUDAExecutionProvider",
+        device: str | torch.device = "cuda",
+        onnx_dtype: torch.dtype = torch.float32,
+        mask_threshold=0.0,
+        max_hole_area=0.0,
+        max_sprinkle_area=0.0,
+        **kwargs,
+    ) -> None:
+        """
+        Uses SAM-2 to compute the image embedding for an image, and then allow mask prediction given prompts.
+
+        Arguments:
+          sam_model (SAM2Base): The model to use for mask prediction.
+          onnx_directory (str): The path of the directory that contains encoder and decoder onnx models.
+          onnx_dtype (torch.dtype): The data type to use for ONNX inputs.
+          mask_threshold (float): The threshold to convert mask logits to binary masks. Default is 0.0.
+          max_hole_area (float): If max_hole_area > 0, we fill small holes in up to
+                                 the maximum area of max_hole_area in low_res_masks.
+          max_sprinkle_area (float): If max_sprinkle_area > 0, we remove small sprinkles up to
+                                     the maximum area of max_sprinkle_area in low_res_masks.
+        """
+        super().__init__(
+            sam_model, mask_threshold=mask_threshold, max_hole_area=max_hole_area, max_sprinkle_area=max_sprinkle_area
+        )
+
+        logger.debug("self.device=%s, device=%s", self.device, device)
+
+        # This model is exported by image_encoder.py.
+        self.encoder_session = create_session(
+            image_encoder_onnx_path,
+            session_options=None,
+            provider=provider,
+            device=device,
+            enable_cuda_graph=False,
+        )
+        self.onnx_dtype = onnx_dtype
+
+        # This model is exported by image_decoder.py. It outputs only one mask.
+        self.decoder_session = create_session(
+            image_decoder_onnx_path,
+            session_options=None,
+            provider=provider,
+            device=device,
+            enable_cuda_graph=False,
+        )
+
+        # This model is exported by image_decoder.py. It outputs multiple (3) masks.
+        self.decoder_session_multi_out = create_session(
+            image_decoder_multi_onnx_path,
+            session_options=None,
+            provider=provider,
+            device=device,
+            enable_cuda_graph=False,
+        )
+
+    @torch.no_grad()
+    def set_image(self, image: np.ndarray | Image):
+        """
+        Calculates the image embeddings for the provided image.
+
+        Arguments:
+          image (np.ndarray or PIL Image): The input image to embed in RGB format.
+              The image should be in HWC format if np.ndarray, or WHC format if PIL Image with pixel values in [0, 255].
+        """
+        self.reset_predictor()
+        # Transform the image to the form expected by the model
+        if isinstance(image, np.ndarray):
+            # For numpy array image, we assume (HxWxC) format.
+            self._orig_hw = [image.shape[:2]]
+        elif isinstance(image, Image):
+            w, h = image.size
+            self._orig_hw = [(h, w)]
+        else:
+            raise NotImplementedError("Image format not supported")
+
+        input_image = self._transforms(image)
+        input_image = input_image[None, ...].to(self.device)
+
+        assert len(input_image.shape) == 4 and input_image.shape[1] == 3, (
+            f"input_image must be of size 1x3xHxW, got {input_image.shape}"
+        )
+
+        # Computing image embeddings for the provided image
+        io_shapes = encoder_shape_dict(batch_size=1, height=input_image.shape[2], width=input_image.shape[3])
+        self.encoder_session.allocate_buffers(io_shapes)
+
+        feed_dict = {"image": input_image.to(self.onnx_dtype).to(self.device)}
+
+        for key, value in feed_dict.items():
+            logger.debug(f"{key}: {value.shape}, {value.dtype}")
+        logger.debug(f"encoder onnx: {self.encoder_session.ort_session._model_path}")
+
+        ort_outputs = self.encoder_session.infer(feed_dict)
+
+        self._features = {
+            "image_embed": ort_outputs["image_embeddings"],
+            "high_res_feats": [ort_outputs[f"image_features_{i}"] for i in range(2)],
+        }
+        self._is_image_set = True
+        logging.info("Image embeddings computed.")
+
+    @torch.no_grad()
+    def _predict(
+        self,
+        point_coords: torch.Tensor | None,
+        point_labels: torch.Tensor | None,
+        boxes: torch.Tensor | None = None,
+        mask_input: torch.Tensor | None = None,
+        multimask_output: bool = True,
+        return_logits: bool = False,
+        img_idx: int = -1,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Predict masks for the given input prompts, using the currently set image.
+        Input prompts are batched torch tensors and are expected to already be
+        transformed to the input frame using SAM2Transforms.
+
+        Arguments:
+          point_coords (torch.Tensor or None): A BxNx2 array of point prompts to the
+            model. Each point is in (X,Y) in pixels.
+          point_labels (torch.Tensor or None): A BxN array of labels for the
+            point prompts. 1 indicates a foreground point and 0 indicates a
+            background point.
+          boxes (np.ndarray or None): A Bx4 array given a box prompt to the
+            model, in XYXY format.
+          mask_input (np.ndarray): A low resolution mask input to the model, typically
+            coming from a previous prediction iteration. Has form Bx1xHxW, where
+            for SAM, H=W=256. Masks returned by a previous iteration of the
+            predict method do not need further transformation.
+          multimask_output (bool): If true, the model will return three masks.
+            For ambiguous input prompts (such as a single click), this will often
+            produce better masks than a single prediction. If only a single
+            mask is needed, the model's predicted quality score can be used
+            to select the best mask. For non-ambiguous prompts, such as multiple
+            input prompts, multimask_output=False can give better results.
+          return_logits (bool): If true, returns un-thresholded masks logits
+            instead of a binary mask.
+
+        Returns:
+          (torch.Tensor): The output masks in BxCxHxW format, where C is the
+            number of masks, and (H, W) is the original image size.
+          (torch.Tensor): An array of shape BxC containing the model's
+            predictions for the quality of each mask.
+          (torch.Tensor): An array of shape BxCxHxW, where C is the number
+            of masks and H=W=256. These low res logits can be passed to
+            a subsequent iteration as mask input.
+        """
+        assert not return_logits  # onnx model is exported for returning bool masks.
+
+        if not self._is_image_set:
+            raise RuntimeError("An image must be set with .set_image(...) before mask prediction.")
+
+        if point_coords is not None:
+            concat_points = (point_coords, point_labels)
+        else:
+            concat_points = None
+
+        # Embed prompts
+        if boxes is not None:
+            box_coords = boxes.reshape(-1, 2, 2)
+            box_labels = torch.tensor([[2, 3]], dtype=torch.int, device=boxes.device)
+            box_labels = box_labels.repeat(boxes.size(0), 1)
+            # we merge "boxes" and "points" into a single "concat_points" input (where
+            # boxes are added at the beginning) to sam_prompt_encoder
+            if concat_points is not None:
+                concat_coords = torch.cat([box_coords, concat_points[0]], dim=1)
+                concat_labels = torch.cat([box_labels, concat_points[1]], dim=1)
+                concat_points = (concat_coords, concat_labels)
+            else:
+                concat_points = (box_coords, box_labels)
+
+        assert concat_points is not None
+        num_labels = concat_points[0].shape[0]
+        shape_dict = decoder_shape_dict(
+            original_image_height=self._orig_hw[img_idx][0],
+            original_image_width=self._orig_hw[img_idx][1],
+            num_labels=num_labels,
+            max_points=concat_points[0].shape[1],
+            num_masks=3 if multimask_output else 1,
+        )
+        if multimask_output:
+            decoder_session = self.decoder_session_multi_out
+        else:
+            decoder_session = self.decoder_session
+
+        decoder_session.allocate_buffers(shape_dict)
+
+        image_features_0 = self._features["high_res_feats"][0][img_idx].unsqueeze(0)
+        image_features_1 = self._features["high_res_feats"][1][img_idx].unsqueeze(0)
+        image_embeddings = self._features["image_embed"][img_idx].unsqueeze(0)
+
+        if mask_input is None:
+            input_masks = torch.zeros(num_labels, 1, 256, 256, dtype=self.onnx_dtype, device=self.device)
+            has_input_masks = torch.zeros(num_labels, dtype=self.onnx_dtype, device=self.device)
+        else:
+            input_masks = mask_input[img_idx].unsqueeze(0).repeat(num_labels, 1, 1, 1)
+            has_input_masks = torch.ones(num_labels, dtype=self.onnx_dtype, device=self.device)
+
+        feed_dict = {
+            "image_embeddings": image_embeddings.contiguous().to(dtype=self.onnx_dtype).to(self.device),
+            "image_features_0": image_features_0.contiguous().to(dtype=self.onnx_dtype).to(self.device),
+            "image_features_1": image_features_1.contiguous().to(dtype=self.onnx_dtype).to(self.device),
+            "point_coords": concat_points[0].to(dtype=self.onnx_dtype).to(self.device),
+            "point_labels": concat_points[1].to(dtype=torch.int32).to(self.device),
+            "input_masks": input_masks.to(dtype=self.onnx_dtype).to(self.device),
+            "has_input_masks": has_input_masks.to(dtype=self.onnx_dtype).to(self.device),
+            "original_image_size": torch.tensor(self._orig_hw[img_idx], dtype=torch.int32, device=self.device),
+        }
+
+        for key, value in feed_dict.items():
+            logger.debug(f"{key}: {value.shape}, {value.dtype}")
+        logger.debug(f"decoder onnx: {self.decoder_session.ort_session._model_path}")
+
+        ort_outputs = decoder_session.infer(feed_dict)
+
+        masks = ort_outputs["masks"]
+        iou_predictions = ort_outputs["iou_predictions"]
+        low_res_masks = ort_outputs["low_res_masks"]
+
+        return torch.Tensor(masks), torch.Tensor(iou_predictions), torch.Tensor(low_res_masks)