Spaces:
Running
on
T4
Running
on
T4
import torch | |
import numpy as np | |
from torchvision.transforms import ToTensor | |
GPU_EFFICIENT_SAM_CHECKPOINT = "efficient_sam_s_gpu.jit" | |
CPU_EFFICIENT_SAM_CHECKPOINT = "efficient_sam_s_cpu.jit" | |
def load(device: torch.device) -> torch.jit.ScriptModule: | |
if device.type == "cuda": | |
model = torch.jit.load(GPU_EFFICIENT_SAM_CHECKPOINT) | |
else: | |
model = torch.jit.load(CPU_EFFICIENT_SAM_CHECKPOINT) | |
model.eval() | |
return model | |
def inference_with_box( | |
image: np.ndarray, | |
box: np.ndarray, | |
model: torch.jit.ScriptModule, | |
device: torch.device | |
) -> np.ndarray: | |
bbox = torch.reshape(torch.tensor(box), [1, 1, 2, 2]) | |
bbox_labels = torch.reshape(torch.tensor([2, 3]), [1, 1, 2]) | |
img_tensor = ToTensor()(image) | |
predicted_logits, predicted_iou = model( | |
img_tensor[None, ...].to(device), | |
bbox.to(device), | |
bbox_labels.to(device), | |
) | |
predicted_logits = predicted_logits.cpu() | |
all_masks = torch.ge(torch.sigmoid(predicted_logits[0, 0, :, :, :]), 0.5).numpy() | |
predicted_iou = predicted_iou[0, 0, ...].cpu().detach().numpy() | |
max_predicted_iou = -1 | |
selected_mask_using_predicted_iou = None | |
for m in range(all_masks.shape[0]): | |
curr_predicted_iou = predicted_iou[m] | |
if ( | |
curr_predicted_iou > max_predicted_iou | |
or selected_mask_using_predicted_iou is None | |
): | |
max_predicted_iou = curr_predicted_iou | |
selected_mask_using_predicted_iou = all_masks[m] | |
return selected_mask_using_predicted_iou | |