diegokauer/segmentation-backend

1import os2import logging3import torch4import cv25import numpy as np6 7from typing import List, Dict, Optional8from label_studio_ml.utils import get_image_local_path, InMemoryLRUDictCache9 10logger = logging.getLogger(__name__)11 12VITH_CHECKPOINT = os.environ.get("VITH_CHECKPOINT")13ONNX_CHECKPOINT = os.environ.get("ONNX_CHECKPOINT")14MOBILESAM_CHECKPOINT = os.environ.get("MOBILESAM_CHECKPOINT", "mobile_sam.pt")15LABEL_STUDIO_ACCESS_TOKEN = os.environ.get("LABEL_STUDIO_ACCESS_TOKEN")16LABEL_STUDIO_HOST = os.environ.get("LABEL_STUDIO_HOST")17 18 19class SAMPredictor(object):20 21    def __init__(self, model_choice):22        self.model_choice = model_choice23 24        # cache for embeddings25        # TODO: currently it supports only one image in cache,26        #   since predictor.set_image() should be called each time the new image comes27        #   before making predictions28        #   to extend it to >1 image, we need to store the "active image" state in the cache29        self.cache = InMemoryLRUDictCache(1)30 31        # if you're not using CUDA, use "cpu" instead .... good luck not burning your computer lol32        self.device = "cuda" if torch.cuda.is_available() else "cpu"33        logger.debug(f"Using device {self.device}")34 35        if model_choice == 'ONNX':36            import onnxruntime37            from segment_anything import sam_model_registry, SamPredictor38 39            self.model_checkpoint = VITH_CHECKPOINT40            if self.model_checkpoint is None:41                raise FileNotFoundError("VITH_CHECKPOINT is not set: please set it to the path to the SAM checkpoint")42            if ONNX_CHECKPOINT is None:43                raise FileNotFoundError("ONNX_CHECKPOINT is not set: please set it to the path to the ONNX checkpoint")44            logger.info(f"Using ONNX checkpoint {ONNX_CHECKPOINT} and SAM checkpoint {self.model_checkpoint}")45 46            self.ort = onnxruntime.InferenceSession(ONNX_CHECKPOINT)47            reg_key = "vit_h"48 49        elif model_choice == 'SAM':50            from segment_anything import SamPredictor, sam_model_registry51 52            self.model_checkpoint = VITH_CHECKPOINT53            if self.model_checkpoint is None:54                raise FileNotFoundError("VITH_CHECKPOINT is not set: please set it to the path to the SAM checkpoint")55 56            logger.info(f"Using SAM checkpoint {self.model_checkpoint}")57            reg_key = "vit_h"58 59        elif model_choice == 'MobileSAM':60            from mobile_sam import SamPredictor, sam_model_registry61 62            self.model_checkpoint = MOBILESAM_CHECKPOINT63            if not self.model_checkpoint:64                raise FileNotFoundError("MOBILE_CHECKPOINT is not set: please set it to the path to the MobileSAM checkpoint")65            logger.info(f"Using MobileSAM checkpoint {self.model_checkpoint}")66            reg_key = 'vit_t'67        else:68            raise ValueError(f"Invalid model choice {model_choice}")69 70        sam = sam_model_registry[reg_key](checkpoint=self.model_checkpoint)71        sam.to(device=self.device)72        self.predictor = SamPredictor(sam)73 74    @property75    def model_name(self):76        return f'{self.model_choice}:{self.model_checkpoint}:{self.device}'77 78    def set_image(self, img_path, calculate_embeddings=True):79        payload = self.cache.get(img_path)80        if payload is None:81            # Get image and embeddings82            logger.debug(f'Payload not found for {img_path} in `IN_MEM_CACHE`: calculating from scratch')83            image_path = get_image_local_path(84                img_path,85                label_studio_access_token=LABEL_STUDIO_ACCESS_TOKEN,86                label_studio_host=LABEL_STUDIO_HOST87            )88            image = cv2.imread(image_path)89            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)90            self.predictor.set_image(image)91            payload = {'image_shape': image.shape[:2]}92            logger.debug(f'Finished set_image({img_path}) in `IN_MEM_CACHE`: image shape {image.shape[:2]}')93            if calculate_embeddings:94                image_embedding = self.predictor.get_image_embedding().cpu().numpy()95                payload['image_embedding'] = image_embedding96                logger.debug(f'Finished storing embeddings for {img_path} in `IN_MEM_CACHE`: '97                             f'embedding shape {image_embedding.shape}')98            self.cache.put(img_path, payload)99        else:100            logger.debug(f"Using embeddings for {img_path} from `IN_MEM_CACHE`")101        return payload102 103    def predict_onnx(104        self,105        img_path,106        point_coords: Optional[List[List]] = None,107        point_labels: Optional[List] = None,108        input_box: Optional[List] = None109    ):110        # calculate embeddings111        payload = self.set_image(img_path, calculate_embeddings=True)112        image_shape = payload['image_shape']113        image_embedding = payload['image_embedding']114 115        onnx_point_coords = np.array(point_coords, dtype=np.float32) if point_coords else None116        onnx_point_labels = np.array(point_labels, dtype=np.float32) if point_labels else None117        onnx_box_coords = np.array(input_box, dtype=np.float32).reshape(2, 2) if input_box else None118 119        onnx_coords, onnx_labels = None, None120        if onnx_point_coords is not None and onnx_box_coords is not None:121            # both keypoints and boxes are present122            onnx_coords = np.concatenate([onnx_point_coords, onnx_box_coords], axis=0)[None, :, :]123            onnx_labels = np.concatenate([onnx_point_labels, np.array([2, 3])], axis=0)[None, :].astype(np.float32)124 125        elif onnx_point_coords is not None:126            # only keypoints are present127            onnx_coords = np.concatenate([onnx_point_coords, np.array([[0.0, 0.0]])], axis=0)[None, :, :]128            onnx_labels = np.concatenate([onnx_point_labels, np.array([-1])], axis=0)[None, :].astype(np.float32)129 130        elif onnx_box_coords is not None:131            # only boxes are present132            raise NotImplementedError("Boxes without keypoints are not supported yet")133 134        onnx_coords = self.predictor.transform.apply_coords(onnx_coords, image_shape).astype(np.float32)135 136        # TODO: support mask inputs137        onnx_mask_input = np.zeros((1, 1, 256, 256), dtype=np.float32)138 139        onnx_has_mask_input = np.zeros(1, dtype=np.float32)140 141        ort_inputs = {142            "image_embeddings": image_embedding,143            "point_coords": onnx_coords,144            "point_labels": onnx_labels,145            "mask_input": onnx_mask_input,146            "has_mask_input": onnx_has_mask_input,147            "orig_im_size": np.array(image_shape, dtype=np.float32)148        }149 150        masks, prob, low_res_logits = self.ort.run(None, ort_inputs)151        masks = masks > self.predictor.model.mask_threshold152        mask = masks[0, 0, :, :].astype(np.uint8)  # each mask has shape [H, W]153        prob = float(prob[0][0])154        # TODO: support the real multimask output as in https://github.com/facebookresearch/segment-anything/blob/main/notebooks/predictor_example.ipynb155        return {156            'masks': [mask],157            'probs': [prob]158        }159 160    def predict_sam(161        self,162        img_path,163        point_coords: Optional[List[List]] = None,164        point_labels: Optional[List] = None,165        input_box: Optional[List] = None166    ):167        self.set_image(img_path, calculate_embeddings=False)168        point_coords = np.array(point_coords, dtype=np.float32) if point_coords else None169        point_labels = np.array(point_labels, dtype=np.float32) if point_labels else None170        input_box = np.array(input_box, dtype=np.float32) if input_box else None171 172        masks, probs, logits = self.predictor.predict(173            point_coords=point_coords,174            point_labels=point_labels,175            box=input_box,176            # TODO: support multimask output177            multimask_output=False178        )179        mask = masks[0, :, :].astype(np.uint8)  # each mask has shape [H, W]180        prob = float(probs[0])181        return {182            'masks': [mask],183            'probs': [prob]184        }185 186    def predict(187        self, img_path: str,188        point_coords: Optional[List[List]] = None,189        point_labels: Optional[List] = None,190        input_box: Optional[List] = None191    ):192        if self.model_choice == 'ONNX':193            return self.predict_onnx(img_path, point_coords, point_labels, input_box)194        elif self.model_choice in ('SAM', 'MobileSAM'):195            return self.predict_sam(img_path, point_coords, point_labels, input_box)196        else:197            raise NotImplementedError(f"Model choice {self.model_choice} is not supported yet")198 199