chore: rework SAM model

label_studio_ml/examples/segment_anything_model/Dockerfile

@@ -1,21 +1,30 @@
-FROM python:3.8-slim
+FROM python:3.8-slim-buster
 
 WORKDIR /app
 
 ENV PYTHONUNBUFFERED=True \
-    PORT=9090
-
+    PORT=9090 \
+    MODEL_DIR=/data/models \
+    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
+
+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/*
+
+# Copy and run the model download script
+COPY download_models.sh .
+RUN bash /app/download_models.sh
+
+# Install Python dependencies
 COPY requirements.txt .
-
-RUN apt-get update && apt-get install -y wget git python3-opencv \
-    && pip install --no-cache-dir -r requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
 
 COPY . ./
 
-# Make the start.sh script executable
-RUN chmod +x /app/start.sh
-
 EXPOSE 9090
 
-# Set the start.sh script as the entrypoint
-ENTRYPOINT ["/app/start.sh"]
+CMD ["/app/start.sh"]

label_studio_ml/examples/segment_anything_model/README.md

@@ -14,21 +14,6 @@ To start the server with lightweight mobile version of SAM, run the following co
 docker-compose up
 ```
 
-### Download models
-
-If you're planning to use the original SAM architecture, you will need to download the model weights. 
-You can do this by running the following command:
-
-```bash
-./download_models.sh
-```
-
-then rebuild the server with the following command:
-
-```bash
-docker-compose up --build
-```
-
 # Intro
 
 There are three models in this repo that you can use.

label_studio_ml/examples/segment_anything_model/docker-compose.yml

@@ -16,10 +16,6 @@ services:
     environment:
       # Change this to your model name
       - SAM_CHOICE=MobileSAM
-      - 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
       - LABEL_STUDIO_HOST=
@@ -28,4 +24,3 @@ services:
       - 9090:9090
     volumes:
       - "./data/server:/data"
-      - "./models:/app/models"

label_studio_ml/examples/segment_anything_model/download_models.sh

@@ -1,13 +1,22 @@
 #!/bin/bash
 
-# Make sure the models directory exists
-mkdir -p models
+MODELS_DIR="models"
+mkdir -p ${MODELS_DIR}
 
-# check if file `sam_vit_h_4b8939.pth` exists, otherwise download the model
-[ -f models/sam_vit_h_4b8939.pth ] ||
-wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth -P models/
+download_model() {
+  FILE_PATH="${MODELS_DIR}/$1"
+  URL="$2"
 
-# check if file `mobile_sam.pt` exists, otherwise download the model
-[ -f models/mobile_sam.pt ] ||
-wget https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt -P models/
+  if [ ! -f "${FILE_PATH}" ]; then
+    wget -q "${URL}" -P ${MODELS_DIR}/
+  fi
+}
 
+# Model files and their corresponding URLs
+declare -A MODELS
+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

label_studio_ml/examples/segment_anything_model/start.sh

@@ -1,22 +1,8 @@
 #!/bin/bash
 
-# check if file `mobile_sam.pt` exists, otherwise download the model
-# we keep this model inside the container, so we don't need to download it upfront
-[ -f "$MOBILESAM_CHECKPOINT" ] ||
-wget https://github.com/ChaoningZhang/MobileSAM/raw/master/weights/mobile_sam.pt -P $(dirname "$MOBILESAM_CHECKPOINT")
-
-# Check if the file specified in the ONNX_CHECKPOINT environment variable exists
-if [ -f "$VITH_CHECKPOINT" ] && [ ! -f "$ONNX_CHECKPOINT" ]; then
-  # Run the python onnxconverter.py script if the file is not found
+# Run converter for ONNX
+if [ "$SAM_CHOICE" == "ONNX" ]; then
   python3 onnxconverter.py
-else
-  # if VITH_CHECKPOINT is not found, print a message to the console
-  if [ ! -f "$VITH_CHECKPOINT" ]; then
-    echo "VITH checkpoint not found in $VITH_CHECKPOINT. Run download_models.sh to download the model."
-  else
-    # Otherwise, print a message to the console
-    echo "ONNX checkpoint found in $ONNX_CHECKPOINT, skipping conversion"
-  fi
 fi
 
 # Execute the gunicorn command