Initial commit of the Label Studio Segment Anything space

Implementation of a Label Studio ML backend using MobileSAM
for image segmentation.

Dockerfile

@@ -0,0 +1,39 @@
+FROM python:3.8-slimjjjjjj
+
+# Install Dependencies
+RUN apt-get update -q \
+ && apt-get install -qy --no-install-recommends wget git libopencv-dev python3-opencv \
+ && apt-get autoremove -y \
+ && apt-get clean \
+ && rm -rf /var/lib/apt/lists/*
+
+# Set up a non-root user
+RUN useradd -m -u 1000 user \
+ && mkdir /app \
+ && chown -R user /app
+
+ # Switch to the "user" user
+USER user
+
+ # Set the working directory to the user's home directory
+WORKDIR /app
+
+ENV PYTHONUNBUFFERED=True \
+    VITH_CHECKPOINT=/app/models/sam_vit_h_4b8939.pth \
+    MOBILESAM_CHECKPOINT=/app/models/mobile_sam.pt \
+    ONNX_CHECKPOINT=/app/models/sam_onnx_quantized_example.onnx \
+    PORT=7860
+
+# Copy and run the model download script
+COPY download_models.sh .
+RUN bash /app/download_models.sh
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip install --user --no-cache-dir -r requirements.txt
+
+COPY . ./
+
+EXPOSE 7860
+
+CMD ["/app/start.sh"]

_wsgi.py

@@ -0,0 +1,113 @@
+import os
+import argparse
+import logging
+import logging.config
+import json
+
+logging.config.dictConfig({
+  "version": 1,
+  "formatters": {
+    "standard": {
+      "format": "[%(asctime)s] [%(levelname)s] [%(name)s::%(funcName)s::%(lineno)d] %(message)s"
+    }
+  },
+  "handlers": {
+    "console": {
+      "class": "logging.StreamHandler",
+      "level": os.getenv('LOG_LEVEL', 'INFO'),
+      "stream": "ext://sys.stdout",
+      "formatter": "standard"
+    }
+  },
+  "root": {
+    "level": os.getenv('LOG_LEVEL', 'INFO'),
+    "handlers": [
+      "console"
+    ],
+    "propagate": True
+  }
+})
+
+from label_studio_ml.api import init_app
+from model import SamMLBackend
+
+_DEFAULT_CONFIG_PATH = os.path.join(os.path.dirname(__file__), 'config.json')
+
+
+def get_kwargs_from_config(config_path=_DEFAULT_CONFIG_PATH):
+    if not os.path.exists(config_path):
+        return dict()
+    with open(config_path) as f:
+        config = json.load(f)
+    assert isinstance(config, dict)
+    return config
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Label studio')
+    parser.add_argument(
+        '-p', '--port', dest='port', type=int, default=9090,
+        help='Server port')
+    parser.add_argument(
+        '--host', dest='host', type=str, default='0.0.0.0',
+        help='Server host')
+    parser.add_argument(
+        '--kwargs', '--with', dest='kwargs', metavar='KEY=VAL', nargs='+', type=lambda kv: kv.split('='),
+        help='Additional LabelStudioMLBase model initialization kwargs')
+    parser.add_argument(
+        '-d', '--debug', dest='debug', action='store_true',
+        help='Switch debug mode')
+    parser.add_argument(
+        '--log-level', dest='log_level', choices=['DEBUG', 'INFO', 'WARNING', 'ERROR'], default=None,
+        help='Logging level')
+    parser.add_argument(
+        '--model-dir', dest='model_dir', default=os.path.dirname(__file__),
+        help='Directory where models are stored (relative to the project directory)')
+    parser.add_argument(
+        '--check', dest='check', action='store_true',
+        help='Validate model instance before launching server')
+
+    args = parser.parse_args()
+
+    # setup logging level
+    if args.log_level:
+        logging.root.setLevel(args.log_level)
+
+    def isfloat(value):
+        try:
+            float(value)
+            return True
+        except ValueError:
+            return False
+
+    def parse_kwargs():
+        param = dict()
+        for k, v in args.kwargs:
+            if v.isdigit():
+                param[k] = int(v)
+            elif v == 'True' or v == 'true':
+                param[k] = True
+            elif v == 'False' or v == 'False':
+                param[k] = False
+            elif isfloat(v):
+                param[k] = float(v)
+            else:
+                param[k] = v
+        return param
+
+    kwargs = get_kwargs_from_config()
+
+    if args.kwargs:
+        kwargs.update(parse_kwargs())
+
+    if args.check:
+        print('Check "' + SamMLBackend.__name__ + '" instance creation..')
+        model = SamMLBackend(**kwargs)
+
+    app = init_app(model_class=SamMLBackend)
+
+    app.run(host=args.host, port=args.port, debug=args.debug)
+
+else:
+    # for uWSGI use
+    app = init_app(model_class=SamMLBackend)

download_models.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+MODELS_DIR="models"
+mkdir -p ${MODELS_DIR}
+
+download_model() {
+  FILE_PATH="${MODELS_DIR}/$1"
+  URL="$2"
+
+  if [ ! -f "${FILE_PATH}" ]; then
+    wget -q "${URL}" -P ${MODELS_DIR}/
+  fi
+}
+
+# Model files and their corresponding URLs
+declare -A MODELS
+# We just run with MobileSAM for this example
+# MODELS["sam_vit_h_4b8939.pth"]="https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
+MODELS["mobile_sam.pt"]="https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt"
+
+for model in "${!MODELS[@]}"; do
+  download_model "${model}" "${MODELS[${model}]}"
+done

model.py

@@ -0,0 +1,145 @@
+import os
+
+from label_studio_converter import brush
+from typing import List, Dict, Optional
+from uuid import uuid4
+from sam_predictor import SAMPredictor
+from label_studio_ml.model import LabelStudioMLBase
+
+SAM_CHOICE = os.environ.get("SAM_CHOICE", "MobileSAM")  # other option is just SAM
+PREDICTOR = SAMPredictor(SAM_CHOICE)
+
+
+class SamMLBackend(LabelStudioMLBase):
+
+    def predict(self, tasks: List[Dict], context: Optional[Dict] = None, **kwargs) -> List[Dict]:
+        """ Returns the predicted mask for a smart keypoint that has been placed."""
+
+        from_name, to_name, value = self.get_first_tag_occurence('BrushLabels', 'Image')
+
+        if not context or not context.get('result'):
+            # if there is no context, no interaction has happened yet
+            return []
+
+        image_width = context['result'][0]['original_width']
+        image_height = context['result'][0]['original_height']
+
+        # collect context information
+        point_coords = []
+        point_labels = []
+        input_box = None
+        selected_label = None
+        for ctx in context['result']:
+            x = ctx['value']['x'] * image_width / 100
+            y = ctx['value']['y'] * image_height / 100
+            ctx_type = ctx['type']
+            selected_label = ctx['value'][ctx_type][0]
+            if ctx_type == 'keypointlabels':
+                point_labels.append(int(ctx['is_positive']))
+                point_coords.append([int(x), int(y)])
+            elif ctx_type == 'rectanglelabels':
+                box_width = ctx['value']['width'] * image_width / 100
+                box_height = ctx['value']['height'] * image_height / 100
+                input_box = [int(x), int(y), int(box_width + x), int(box_height + y)]
+
+        print(f'Point coords are {point_coords}, point labels are {point_labels}, input box is {input_box}')
+
+        img_path = tasks[0]['data'][value]
+        predictor_results = PREDICTOR.predict(
+            img_path=img_path,
+            point_coords=point_coords or None,
+            point_labels=point_labels or None,
+            input_box=input_box
+        )
+
+        predictions = self.get_results(
+            masks=predictor_results['masks'],
+            probs=predictor_results['probs'],
+            width=image_width,
+            height=image_height,
+            from_name=from_name,
+            to_name=to_name,
+            label=selected_label)
+
+        return predictions
+
+    def get_results(self, masks, probs, width, height, from_name, to_name, label):
+        results = []
+        for mask, prob in zip(masks, probs):
+            # creates a random ID for your label everytime so no chance for errors
+            label_id = str(uuid4())[:4]
+            # converting the mask from the model to RLE format which is usable in Label Studio
+            mask = mask * 255
+            rle = brush.mask2rle(mask)
+
+            results.append({
+                'id': label_id,
+                'from_name': from_name,
+                'to_name': to_name,
+                'original_width': width,
+                'original_height': height,
+                'image_rotation': 0,
+                'value': {
+                    'format': 'rle',
+                    'rle': rle,
+                    'brushlabels': [label],
+                },
+                'score': prob,
+                'type': 'brushlabels',
+                'readonly': False
+            })
+
+        return [{
+            'result': results,
+            'model_version': PREDICTOR.model_name
+        }]
+
+
+if __name__ == '__main__':
+    # test the model
+    model = SamMLBackend()
+    model.use_label_config('''
+    <View>
+        <Image name="image" value="$image" zoom="true"/>
+        <BrushLabels name="tag" toName="image">
+            <Label value="Banana" background="#FF0000"/>
+            <Label value="Orange" background="#0d14d3"/>
+        </BrushLabels>
+        <KeyPointLabels name="tag2" toName="image" smart="true" >
+            <Label value="Banana" background="#000000" showInline="true"/>
+            <Label value="Orange" background="#000000" showInline="true"/>
+        </KeyPointLabels>
+        <RectangleLabels name="tag3" toName="image"  >
+            <Label value="Banana" background="#000000" showInline="true"/>
+            <Label value="Orange" background="#000000" showInline="true"/>
+        </RectangleLabels>
+    </View>
+    ''')
+    results = model.predict(
+        tasks=[{
+            'data': {
+                'image': 'https://s3.amazonaws.com/htx-pub/datasets/images/125245483_152578129892066_7843809718842085333_n.jpg'
+            }}],
+        context={
+            'result': [{
+                'original_width': 1080,
+                'original_height': 1080,
+                'image_rotation': 0,
+                'value': {
+                    'x': 49.441786283891545,
+                    'y': 59.96810207336522,
+                    'width': 0.3189792663476874,
+                    'labels': ['Banana'],
+                    'keypointlabels': ['Banana']
+                },
+                'is_positive': True,
+                'id': 'fBWv1t0S2L',
+                'from_name': 'tag2',
+                'to_name': 'image',
+                'type': 'keypointlabels',
+                'origin': 'manual'
+            }]}
+    )
+    import json
+    results[0]['result'][0]['value']['rle'] = f'...{len(results[0]["result"][0]["value"]["rle"])} integers...'
+    print(json.dumps(results, indent=2))

requirements.txt

@@ -0,0 +1,13 @@
+label_studio_converter
+opencv-python
+onnxruntime
+onnx
+torch==2.0.1
+torchvision==0.15.2
+gunicorn==20.1.0
+rq==1.10.1
+timm==0.4.12
+
+segment_anything @ git+https://github.com/facebookresearch/segment-anything.git
+mobile-sam @ git+https://github.com/ChaoningZhang/MobileSAM.git
+label-studio-ml @ git+https://github.com/heartexlabs/label-studio-ml-backend.git

sam_predictor.py

@@ -0,0 +1,198 @@
+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")
+

start.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+# Execute the gunicorn command
+exec /home/user/.local/bin/gunicorn --preload --bind :$PORT --workers 1 --threads 8 --timeout 0 _wsgi:app