Spaces:
Running
Running
File size: 8,666 Bytes
6307f85 |
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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 |
import os
import logging
import torch
import cv2
import numpy as np
from typing import List, Dict, Optional
from label_studio_ml.utils import get_image_local_path, InMemoryLRUDictCache
logger = logging.getLogger(__name__)
VITH_CHECKPOINT = os.environ.get("VITH_CHECKPOINT")
ONNX_CHECKPOINT = os.environ.get("ONNX_CHECKPOINT")
MOBILESAM_CHECKPOINT = os.environ.get("MOBILESAM_CHECKPOINT", "mobile_sam.pt")
LABEL_STUDIO_ACCESS_TOKEN = os.environ.get("LABEL_STUDIO_ACCESS_TOKEN")
LABEL_STUDIO_HOST = os.environ.get("LABEL_STUDIO_HOST")
class SAMPredictor(object):
def __init__(self, model_choice):
self.model_choice = model_choice
# cache for embeddings
# TODO: currently it supports only one image in cache,
# since predictor.set_image() should be called each time the new image comes
# before making predictions
# to extend it to >1 image, we need to store the "active image" state in the cache
self.cache = InMemoryLRUDictCache(1)
# if you're not using CUDA, use "cpu" instead .... good luck not burning your computer lol
self.device = "cuda" if torch.cuda.is_available() else "cpu"
logger.debug(f"Using device {self.device}")
if model_choice == 'ONNX':
import onnxruntime
from segment_anything import sam_model_registry, SamPredictor
self.model_checkpoint = VITH_CHECKPOINT
if self.model_checkpoint is None:
raise FileNotFoundError("VITH_CHECKPOINT is not set: please set it to the path to the SAM checkpoint")
if ONNX_CHECKPOINT is None:
raise FileNotFoundError("ONNX_CHECKPOINT is not set: please set it to the path to the ONNX checkpoint")
logger.info(f"Using ONNX checkpoint {ONNX_CHECKPOINT} and SAM checkpoint {self.model_checkpoint}")
self.ort = onnxruntime.InferenceSession(ONNX_CHECKPOINT)
reg_key = "vit_h"
elif model_choice == 'SAM':
from segment_anything import SamPredictor, sam_model_registry
self.model_checkpoint = VITH_CHECKPOINT
if self.model_checkpoint is None:
raise FileNotFoundError("VITH_CHECKPOINT is not set: please set it to the path to the SAM checkpoint")
logger.info(f"Using SAM checkpoint {self.model_checkpoint}")
reg_key = "vit_h"
elif model_choice == 'MobileSAM':
from mobile_sam import SamPredictor, sam_model_registry
self.model_checkpoint = MOBILESAM_CHECKPOINT
if not self.model_checkpoint:
raise FileNotFoundError("MOBILE_CHECKPOINT is not set: please set it to the path to the MobileSAM checkpoint")
logger.info(f"Using MobileSAM checkpoint {self.model_checkpoint}")
reg_key = 'vit_t'
else:
raise ValueError(f"Invalid model choice {model_choice}")
sam = sam_model_registry[reg_key](checkpoint=self.model_checkpoint)
sam.to(device=self.device)
self.predictor = SamPredictor(sam)
@property
def model_name(self):
return f'{self.model_choice}:{self.model_checkpoint}:{self.device}'
def set_image(self, img_path, calculate_embeddings=True):
payload = self.cache.get(img_path)
if payload is None:
# Get image and embeddings
logger.debug(f'Payload not found for {img_path} in `IN_MEM_CACHE`: calculating from scratch')
image_path = get_image_local_path(
img_path,
label_studio_access_token=LABEL_STUDIO_ACCESS_TOKEN,
label_studio_host=LABEL_STUDIO_HOST
)
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
self.predictor.set_image(image)
payload = {'image_shape': image.shape[:2]}
logger.debug(f'Finished set_image({img_path}) in `IN_MEM_CACHE`: image shape {image.shape[:2]}')
if calculate_embeddings:
image_embedding = self.predictor.get_image_embedding().cpu().numpy()
payload['image_embedding'] = image_embedding
logger.debug(f'Finished storing embeddings for {img_path} in `IN_MEM_CACHE`: '
f'embedding shape {image_embedding.shape}')
self.cache.put(img_path, payload)
else:
logger.debug(f"Using embeddings for {img_path} from `IN_MEM_CACHE`")
return payload
def predict_onnx(
self,
img_path,
point_coords: Optional[List[List]] = None,
point_labels: Optional[List] = None,
input_box: Optional[List] = None
):
# calculate embeddings
payload = self.set_image(img_path, calculate_embeddings=True)
image_shape = payload['image_shape']
image_embedding = payload['image_embedding']
onnx_point_coords = np.array(point_coords, dtype=np.float32) if point_coords else None
onnx_point_labels = np.array(point_labels, dtype=np.float32) if point_labels else None
onnx_box_coords = np.array(input_box, dtype=np.float32).reshape(2, 2) if input_box else None
onnx_coords, onnx_labels = None, None
if onnx_point_coords is not None and onnx_box_coords is not None:
# both keypoints and boxes are present
onnx_coords = np.concatenate([onnx_point_coords, onnx_box_coords], axis=0)[None, :, :]
onnx_labels = np.concatenate([onnx_point_labels, np.array([2, 3])], axis=0)[None, :].astype(np.float32)
elif onnx_point_coords is not None:
# only keypoints are present
onnx_coords = np.concatenate([onnx_point_coords, np.array([[0.0, 0.0]])], axis=0)[None, :, :]
onnx_labels = np.concatenate([onnx_point_labels, np.array([-1])], axis=0)[None, :].astype(np.float32)
elif onnx_box_coords is not None:
# only boxes are present
raise NotImplementedError("Boxes without keypoints are not supported yet")
onnx_coords = self.predictor.transform.apply_coords(onnx_coords, image_shape).astype(np.float32)
# TODO: support mask inputs
onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)
onnx_has_mask_input = np.zeros(1, dtype=np.float32)
ort_inputs = {
"image_embeddings": image_embedding,
"point_coords": onnx_coords,
"point_labels": onnx_labels,
"mask_input": onnx_mask_input,
"has_mask_input": onnx_has_mask_input,
"orig_im_size": np.array(image_shape, dtype=np.float32)
}
masks, prob, low_res_logits = self.ort.run(None, ort_inputs)
masks = masks > self.predictor.model.mask_threshold
mask = masks[0, 0, :, :].astype(np.uint8) # each mask has shape [H, W]
prob = float(prob[0][0])
# TODO: support the real multimask output as in https://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb
return {
'masks': [mask],
'probs': [prob]
}
def predict_sam(
self,
img_path,
point_coords: Optional[List[List]] = None,
point_labels: Optional[List] = None,
input_box: Optional[List] = None
):
self.set_image(img_path, calculate_embeddings=False)
point_coords = np.array(point_coords, dtype=np.float32) if point_coords else None
point_labels = np.array(point_labels, dtype=np.float32) if point_labels else None
input_box = np.array(input_box, dtype=np.float32) if input_box else None
masks, probs, logits = self.predictor.predict(
point_coords=point_coords,
point_labels=point_labels,
box=input_box,
# TODO: support multimask output
multimask_output=False
)
mask = masks[0, :, :].astype(np.uint8) # each mask has shape [H, W]
prob = float(probs[0])
return {
'masks': [mask],
'probs': [prob]
}
def predict(
self, img_path: str,
point_coords: Optional[List[List]] = None,
point_labels: Optional[List] = None,
input_box: Optional[List] = None
):
if self.model_choice == 'ONNX':
return self.predict_onnx(img_path, point_coords, point_labels, input_box)
elif self.model_choice in ('SAM', 'MobileSAM'):
return self.predict_sam(img_path, point_coords, point_labels, input_box)
else:
raise NotImplementedError(f"Model choice {self.model_choice} is not supported yet")
|