Fix dockerfile

label_studio_ml/examples/llm_interactive/model.py

@@ -170,7 +170,7 @@ def predict(self, tasks, context, **kwargs):
                     'id': str(uuid4())[:4],
                     'from_name': textarea_from_name,
                     'to_name': prompt_to_name,
-                    'type': 'choices',
+                    'type': 'textarea',
                     'value': {
                         'text': [response]
                     }

label_studio_ml/examples/segment_anything_model/.dockerignore

@@ -0,0 +1,19 @@
+Dockerfile
+README.md
+*.pyc
+*.pyo
+*.pyd
+__pycache__
+.pytest_cache
+.idea
+docker-compose.yml
+data
+logs
+*.db
+*.sqlite3
+*.sqlite
+*.sql
+*.pt
+*.pth
+*.onnx
+models/*

label_studio_ml/examples/segment_anything_model/Dockerfile

@@ -10,8 +10,15 @@ COPY requirements.txt .
 RUN apt-get update && apt-get install -y wget git python3-opencv \
     && pip install --no-cache-dir -r requirements.txt
 
+# Copy the start.sh script into the container
+COPY start.sh /app/start.sh
+
+# Make the start.sh script executable
+RUN chmod +x /app/start.sh
+
 COPY . ./
 
 EXPOSE 9090
 
-CMD exec gunicorn --preload --bind :$PORT --workers 1 --threads 8 --timeout 0 _wsgi:app
+# Set the start.sh script as the entrypoint
+ENTRYPOINT ["/app/start.sh"]

label_studio_ml/examples/segment_anything_model/_wsgi.py

@@ -29,7 +29,7 @@
 })
 
 from label_studio_ml.api import init_app
-from model import SamModel
+from model import SamMLBackend
 
 _DEFAULT_CONFIG_PATH = os.path.join(os.path.dirname(__file__), 'config.json')
 
@@ -101,13 +101,13 @@ def parse_kwargs():
         kwargs.update(parse_kwargs())
 
     if args.check:
-        print('Check "' + SamModel.__name__ + '" instance creation..')
-        model = SamModel(**kwargs)
+        print('Check "' + SamMLBackend.__name__ + '" instance creation..')
+        model = SamMLBackend(**kwargs)
 
-    app = init_app(model_class=SamModel)
+    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=SamModel)
+    app = init_app(model_class=SamMLBackend)

label_studio_ml/examples/segment_anything_model/docker-compose.yml

@@ -15,9 +15,9 @@ services:
     environment:
       # Change this to your model name
       - SAM_CHOICE=MobileSAM
-      - VITH_CHECKPOINT=/app/sam_vit_h_4b8939.pth
-      - MOBILESAM_CHECKPOINT=/app/mobile_sam.pth
-      - ONNX_CHECKPOINT=/app/sam_onnx_quantized_example.onnx
+      - 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
       - MODEL_DIR=/data/models
       - LOG_LEVEL=DEBUG
       # Add these variables if you want to access the images stored in Label Studio
@@ -27,8 +27,4 @@ services:
       - 9090:9090
     volumes:
       - "./data/server:/data"
-      # You can comment checkpoints you don't need
-      - "./sam_vit_h_4b8939.pth:/app/sam_vit_h_4b8939.pth"
-      - "./mobile_sam.pt:/app/mobile_sam.pt"
-      - "./sam_onnx_example.onnx:/app/sam_onnx_example.onnx"
-      - "./sam_onnx_quantized_example.onnx:/app/sam_onnx_quantized_example.onnx"
+      - "./models:/app/models"

label_studio_ml/examples/segment_anything_model/download_models.sh

@@ -1,10 +1,13 @@
+#!/bin/bash
+
+# Make sure the models directory exists
+mkdir -p models
+
 # check if file `sam_vit_h_4b8939.pth` exists, otherwise download the model
-[ -f sam_vit_h_4b8939.pth ] ||
-wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
+[ -f models/sam_vit_h_4b8939.pth ] ||
+wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth -P models/
 
 # check if file `mobile_sam.pt` exists, otherwise download the model
-[ -f mobile_sam.pt ] ||
-wget https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt
+[ -f models/mobile_sam.pt ] ||
+wget https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt -P models/
 
-# run ONNX conversion
-python onnxconverter.py

label_studio_ml/examples/segment_anything_model/model.py

@@ -10,14 +10,14 @@
 PREDICTOR = SAMPredictor(SAM_CHOICE)
 
 
-class SamModel(LabelStudioMLBase):
+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:
+        if not context or not context.get('result'):
             # if there is no context, no interaction has happened yet
             return []
 
@@ -97,7 +97,7 @@ def get_results(self, masks, probs, width, height, from_name, to_name, label):
 
 if __name__ == '__main__':
     # test the model
-    model = SamModel()
+    model = SamMLBackend()
     model.use_label_config('''
     <View>
         <Image name="image" value="$image" zoom="true"/>

label_studio_ml/examples/segment_anything_model/sam_predictor.py

@@ -11,7 +11,7 @@
 
 VITH_CHECKPOINT = os.environ.get("VITH_CHECKPOINT", "sam_vit_h_4b8939.pth")
 ONNX_CHECKPOINT = os.environ.get("ONNX_CHECKPOINT", "sam_onnx_quantized_example.onnx")
-MOBILE_CHECKPOINT = os.environ.get("MOBILE_CHECKPOINT", "mobile_sam.pt")
+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")
 
@@ -59,7 +59,7 @@ def __init__(self, model_choice):
         elif model_choice == 'MobileSAM':
             from mobile_sam import SamPredictor, sam_model_registry
 
-            self.model_checkpoint = MOBILE_CHECKPOINT
+            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}")

label_studio_ml/examples/segment_anything_model/start.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+
+# Check if the file specified in the ONNX_CHECKPOINT environment variable exists
+if [ ! -f "$ONNX_CHECKPOINT" ]; then
+  # Run the python onnxconverter.py script if the file is not found
+  python onnxconverter.py
+else
+  # Otherwise, print a message to the console
+  echo "ONNX checkpoint found in $ONNX_CHECKPOINT, skipping conversion"
+fi
+
+# Execute the gunicorn command
+exec gunicorn --preload --bind :$PORT --workers 1 --threads 8 --timeout 0 _wsgi:app