KenjieDec's picture
Update to latest version + sam support?
c8f8b0e verified
import os
from copy import deepcopy
from typing import Dict, List, Tuple
import cv2
import numpy as np
import onnxruntime as ort
import pooch
from jsonschema import validate
from PIL import Image
from PIL.Image import Image as PILImage
from .base import BaseSession
def get_preprocess_shape(oldh: int, oldw: int, long_side_length: int):
scale = long_side_length * 1.0 / max(oldh, oldw)
newh, neww = oldh * scale, oldw * scale
neww = int(neww + 0.5)
newh = int(newh + 0.5)
return (newh, neww)
def apply_coords(coords: np.ndarray, original_size, target_length):
old_h, old_w = original_size
new_h, new_w = get_preprocess_shape(
original_size[0], original_size[1], target_length
)
coords = deepcopy(coords).astype(float)
coords[..., 0] = coords[..., 0] * (new_w / old_w)
coords[..., 1] = coords[..., 1] * (new_h / old_h)
return coords
def get_input_points(prompt):
points = []
labels = []
for mark in prompt:
if mark["type"] == "point":
points.append(mark["data"])
labels.append(mark["label"])
elif mark["type"] == "rectangle":
points.append([mark["data"][0], mark["data"][1]])
points.append([mark["data"][2], mark["data"][3]])
labels.append(2)
labels.append(3)
points, labels = np.array(points), np.array(labels)
return points, labels
def transform_masks(masks, original_size, transform_matrix):
output_masks = []
for batch in range(masks.shape[0]):
batch_masks = []
for mask_id in range(masks.shape[1]):
mask = masks[batch, mask_id]
mask = cv2.warpAffine(
mask,
transform_matrix[:2],
(original_size[1], original_size[0]),
flags=cv2.INTER_LINEAR,
)
batch_masks.append(mask)
output_masks.append(batch_masks)
return np.array(output_masks)
class SamSession(BaseSession):
"""
This class represents a session for the Sam model.
Args:
model_name (str): The name of the model.
sess_opts (ort.SessionOptions): The session options.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
"""
def __init__(
self,
model_name: str,
sess_opts: ort.SessionOptions,
providers=None,
*args,
**kwargs,
):
"""
Initialize a new SamSession with the given model name and session options.
Args:
model_name (str): The name of the model.
sess_opts (ort.SessionOptions): The session options.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
"""
self.model_name = model_name
valid_providers = []
available_providers = ort.get_available_providers()
for provider in providers or []:
if provider in available_providers:
valid_providers.append(provider)
else:
valid_providers.extend(available_providers)
paths = self.__class__.download_models(*args, **kwargs)
self.encoder = ort.InferenceSession(
str(paths[0]),
providers=valid_providers,
sess_options=sess_opts,
)
self.decoder = ort.InferenceSession(
str(paths[1]),
providers=valid_providers,
sess_options=sess_opts,
)
def predict(
self,
img: PILImage,
*args,
**kwargs,
) -> List[PILImage]:
"""
Predict masks for an input image.
This function takes an image as input and performs various preprocessing steps on the image. It then runs the image through an encoder to obtain an image embedding. The function also takes input labels and points as additional arguments. It concatenates the input points and labels with padding and transforms them. It creates an empty mask input and an indicator for no mask. The function then passes the image embedding, point coordinates, point labels, mask input, and has mask input to a decoder. The decoder generates masks based on the input and returns them as a list of images.
Parameters:
img (PILImage): The input image.
*args: Additional arguments.
**kwargs: Additional keyword arguments.
Returns:
List[PILImage]: A list of masks generated by the decoder.
"""
prompt = kwargs.get("sam_prompt", "{}")
schema = {
"type": "array",
"items": {
"type": "object",
"properties": {
"type": {"type": "string"},
"label": {"type": "integer"},
"data": {
"type": "array",
"items": {"type": "number"},
},
},
},
}
validate(instance=prompt, schema=schema)
target_size = 1024
input_size = (684, 1024)
encoder_input_name = self.encoder.get_inputs()[0].name
img = img.convert("RGB")
cv_image = np.array(img)
original_size = cv_image.shape[:2]
scale_x = input_size[1] / cv_image.shape[1]
scale_y = input_size[0] / cv_image.shape[0]
scale = min(scale_x, scale_y)
transform_matrix = np.array(
[
[scale, 0, 0],
[0, scale, 0],
[0, 0, 1],
]
)
cv_image = cv2.warpAffine(
cv_image,
transform_matrix[:2],
(input_size[1], input_size[0]),
flags=cv2.INTER_LINEAR,
)
## encoder
encoder_inputs = {
encoder_input_name: cv_image.astype(np.float32),
}
encoder_output = self.encoder.run(None, encoder_inputs)
image_embedding = encoder_output[0]
embedding = {
"image_embedding": image_embedding,
"original_size": original_size,
"transform_matrix": transform_matrix,
}
## decoder
input_points, input_labels = get_input_points(prompt)
onnx_coord = np.concatenate([input_points, np.array([[0.0, 0.0]])], axis=0)[
None, :, :
]
onnx_label = np.concatenate([input_labels, np.array([-1])], axis=0)[
None, :
].astype(np.float32)
onnx_coord = apply_coords(onnx_coord, input_size, target_size).astype(
np.float32
)
onnx_coord = np.concatenate(
[
onnx_coord,
np.ones((1, onnx_coord.shape[1], 1), dtype=np.float32),
],
axis=2,
)
onnx_coord = np.matmul(onnx_coord, transform_matrix.T)
onnx_coord = onnx_coord[:, :, :2].astype(np.float32)
onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
onnx_has_mask_input = np.zeros(1, dtype=np.float32)
decoder_inputs = {
"image_embeddings": image_embedding,
"point_coords": onnx_coord,
"point_labels": onnx_label,
"mask_input": onnx_mask_input,
"has_mask_input": onnx_has_mask_input,
"orig_im_size": np.array(input_size, dtype=np.float32),
}
masks, _, _ = self.decoder.run(None, decoder_inputs)
inv_transform_matrix = np.linalg.inv(transform_matrix)
masks = transform_masks(masks, original_size, inv_transform_matrix)
mask = np.zeros((masks.shape[2], masks.shape[3], 3), dtype=np.uint8)
for m in masks[0, :, :, :]:
mask[m > 0.0] = [255, 255, 255]
return [Image.fromarray(mask).convert("L")]
@classmethod
def download_models(cls, *args, **kwargs):
"""
Class method to download ONNX model files.
This method is responsible for downloading two ONNX model files from specified URLs and saving them locally. The downloaded files are saved with the naming convention 'name_encoder.onnx' and 'name_decoder.onnx', where 'name' is the value returned by the 'name' method.
Parameters:
cls: The class object.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
tuple: A tuple containing the file paths of the downloaded encoder and decoder models.
"""
model_name = kwargs.get("sam_model", "sam_vit_b_01ec64")
quant = kwargs.get("sam_quant", False)
fname_encoder = f"{model_name}.encoder.onnx"
fname_decoder = f"{model_name}.decoder.onnx"
if quant:
fname_encoder = f"{model_name}.encoder.quant.onnx"
fname_decoder = f"{model_name}.decoder.quant.onnx"
pooch.retrieve(
f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/{fname_encoder}",
None,
fname=fname_encoder,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
pooch.retrieve(
f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/{fname_decoder}",
None,
fname=fname_decoder,
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
if fname_encoder == "sam_vit_h_4b8939.encoder.onnx" and not os.path.exists(
os.path.join(
cls.u2net_home(*args, **kwargs), "sam_vit_h_4b8939.encoder_data.bin"
)
):
content = bytearray()
for i in range(1, 4):
pooch.retrieve(
f"https://github.com/danielgatis/rembg/releases/download/v0.0.0/sam_vit_h_4b8939.encoder_data.{i}.bin",
None,
fname=f"sam_vit_h_4b8939.encoder_data.{i}.bin",
path=cls.u2net_home(*args, **kwargs),
progressbar=True,
)
fbin = os.path.join(
cls.u2net_home(*args, **kwargs),
f"sam_vit_h_4b8939.encoder_data.{i}.bin",
)
content.extend(open(fbin, "rb").read())
os.remove(fbin)
with open(
os.path.join(
cls.u2net_home(*args, **kwargs),
"sam_vit_h_4b8939.encoder_data.bin",
),
"wb",
) as fp:
fp.write(content)
return (
os.path.join(cls.u2net_home(*args, **kwargs), fname_encoder),
os.path.join(cls.u2net_home(*args, **kwargs), fname_decoder),
)
@classmethod
def name(cls, *args, **kwargs):
"""
Class method to return a string value.
This method returns the string value 'sam'.
Parameters:
cls: The class object.
*args: Variable length argument list.
**kwargs: Arbitrary keyword arguments.
Returns:
str: The string value 'sam'.
"""
return "sam"