Spaces:
Runtime error
Runtime error
File size: 3,985 Bytes
69420c9 65665c1 69420c9 65665c1 69420c9 65665c1 435ddbc 69420c9 65665c1 69420c9 65665c1 69420c9 435ddbc 69420c9 65665c1 69420c9 65665c1 69420c9 65665c1 69420c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 |
# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import numpy as np
import cv2
import onnxruntime
from typing import Optional, Tuple
from ..utils.transforms import ResizeLongestSide
class SamPredictorONNX:
mask_threshold: float = 0.0
image_format: str = "RGB"
img_size = 1024
pixel_mean = np.array([123.675, 116.28, 103.53])[None, :, None, None]
pixel_std = np.array([58.395, 57.12, 57.375])[None, :, None, None]
def __init__(
self,
encoder_path: str,
decoder_path: str
) -> None:
super().__init__()
self.encoder = onnxruntime.InferenceSession(encoder_path)
self.decoder = onnxruntime.InferenceSession(decoder_path)
# Set the execution provider to GPU if available
if 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
self.encoder.set_providers(['CUDAExecutionProvider'])
self.decoder.set_providers(['CUDAExecutionProvider'])
self.transform = ResizeLongestSide(self.img_size)
self.reset_image()
def set_image(
self,
image: np.ndarray,
image_format: str = "RGB",
) -> None:
assert image_format in [
"RGB",
"BGR",
], f"image_format must be in ['RGB', 'BGR'], is {image_format}."
if image_format != self.image_format:
image = image[..., ::-1]
# Transform the image to the form expected by the model
input_image = self.transform.apply_image(image)
input_image = input_image.transpose(2, 0, 1)[None, :, :, :]
self.reset_image()
original_size = image.shape[:2]
input_size = tuple(input_image.shape[-2:])
input_image = self.preprocess(input_image).astype(np.float32)
outputs = self.encoder.run(None, {'image': input_image})
features = outputs[0]
return features, input_size, original_size
def predict(
self,
features: np.ndarray,
input_size: Tuple[int, int],
original_size: Tuple[int, int],
point_coords: Optional[np.ndarray] = None,
point_labels: Optional[np.ndarray] = None,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
point_coords = self.transform.apply_coords(point_coords, original_size)
outputs = self.decoder.run(None, {
'image_embeddings': features,
'point_coords': point_coords.astype(np.float32),
'point_labels': point_labels.astype(np.float32)
})
scores, low_res_masks = outputs[0], outputs[1]
masks = self.postprocess_masks(low_res_masks, input_size, original_size)
masks = masks > self.mask_threshold
return masks, scores, low_res_masks
def reset_image(self) -> None:
"""Resets the currently set image."""
self.is_image_set = False
self.features = None
self.orig_h = None
self.orig_w = None
self.input_h = None
self.input_w = None
def preprocess(self, x: np.ndarray):
x = (x - self.pixel_mean) / self.pixel_std
h, w = x.shape[-2:]
padh = self.img_size - h
padw = self.img_size - w
x = np.pad(x, ((0, 0), (0, 0), (0, padh), (0, padw)), mode='constant', constant_values=0)
return x
def postprocess_masks(self, mask: np.ndarray, input_size: Tuple[int, int], original_size: Tuple[int, int]) -> np.ndarray:
mask = mask.squeeze(0).transpose(1, 2, 0)
mask = cv2.resize(mask, (self.img_size, self.img_size), interpolation=cv2.INTER_LINEAR)
mask = mask[:input_size[0], :input_size[1], :]
mask = cv2.resize(mask, (original_size[1], original_size[0]), interpolation=cv2.INTER_LINEAR)
mask = mask.transpose(2, 0, 1)[None, :, :, :]
return mask
|