Update segment-anything-2-video to SAM 2.1

label_studio_ml/examples/segment_anything_2_video/Dockerfile

@@ -1,4 +1,4 @@
-FROM pytorch/pytorch:2.1.2-cuda12.1-cudnn8-runtime
+FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime
 ARG DEBIAN_FRONTEND=noninteractive
 ARG TEST_ENV
 
@@ -22,10 +22,7 @@ ENV PYTHONUNBUFFERED=1 \
     WORKERS=2 \
     THREADS=4 \
     CUDA_HOME=/usr/local/cuda \
-    SEGMENT_ANYTHING_2_REPO_PATH=/segment-anything-2
-
-# this image is already built with cuda 12.1.1, so we don't need to install it again
-# RUN conda install -c "nvidia/label/cuda-12.1.1" cuda -y 
+    PYTHONPATH=/sam2:${PYTHONPATH}
 
 ENV CUDA_HOME=/opt/conda \
     TORCH_CUDA_ARCH_LIST="6.0;6.1;7.0;7.5;8.0;8.6+PTX;8.9;9.0"
@@ -39,16 +36,17 @@ COPY requirements.txt .
 RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
     pip3 install -r requirements.txt
 
-# install segment-anything-2
-RUN cd / && git clone --depth 1 --branch main --single-branch https://github.com/facebookresearch/segment-anything-2.git
-WORKDIR /segment-anything-2
+# install SAM 2.1 (sam2 repo)
+ARG SAM2_REF=main
+RUN cd / && git clone --depth 1 --branch ${SAM2_REF} --single-branch https://github.com/facebookresearch/sam2.git || \
+    (cd / && git clone --depth 1 --branch main --single-branch https://github.com/facebookresearch/sam2.git)
+WORKDIR /sam2
 RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
     pip3 install -e .
 RUN cd checkpoints && ./download_ckpts.sh
 
 WORKDIR /app
 
-# install test requirements if needed
 COPY requirements-test.txt .
 # build only when TEST_ENV="true"
 RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
@@ -57,5 +55,7 @@ RUN --mount=type=cache,target=${PIP_CACHE_DIR},sharing=locked \
     fi
 
 COPY . ./
+# Normalize line endings and ensure executable for start script (fix Windows CRLF issues)
+RUN sed -i 's/\r$//' /app/start.sh && chmod +x /app/start.sh
 
 CMD ["/app/start.sh"]

label_studio_ml/examples/segment_anything_2_video/docker-compose.yml

@@ -1,13 +1,13 @@
-version: "3.8"
-
 services:
   segment_anything_2_video:
     container_name: segment_anything_2_video
     image: humansignal/segment_anything_2_video:v0
     build:
       context: .
       args:
-        TEST_ENV: ${TEST_ENV}
+        TEST_ENV: ${TEST_ENV:-}
+        BASE_IMAGE: ${BASE_IMAGE:-}
+        SAM2_REF: ${SAM2_REF:-main}
     environment:
       # specify these parameters if you want to use basic auth for the model server
       - BASIC_AUTH_USER=
@@ -22,11 +22,11 @@ services:
       # specify the model directory (likely you don't need to change this)
       - MODEL_DIR=/data/models
       # specify device
-      - DEVICE=cuda  # or 'cpu' (coming soon)
+      - DEVICE=cuda
       # SAM2 model config
-      - MODEL_CONFIG=sam2_hiera_l.yaml
+      - MODEL_CONFIG=configs/sam2.1/sam2.1_hiera_t.yaml
       # SAM2 checkpoint
-      - MODEL_CHECKPOINT=sam2_hiera_large.pt
+      - MODEL_CHECKPOINT=sam2.1_hiera_tiny.pt
 
       # Specify the Label Studio URL and API key to access
       # uploaded, local storage and cloud storage files.
@@ -35,7 +35,15 @@ services:
       # Determine the actual IP using 'ifconfig' (Linux/Mac) or 'ipconfig' (Windows).
       - LABEL_STUDIO_URL=
       - LABEL_STUDIO_API_KEY=
+      - MAX_FRAMES_TO_TRACK=10
     ports:
       - "9090:9090"
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
     volumes:
       - "./data/server:/data"

label_studio_ml/examples/segment_anything_2_video/model.py

@@ -19,7 +19,6 @@
 
 
 DEVICE = os.getenv('DEVICE', 'cuda')
-SEGMENT_ANYTHING_2_REPO_PATH = os.getenv('SEGMENT_ANYTHING_2_REPO_PATH', 'segment-anything-2')
 MODEL_CONFIG = os.getenv('MODEL_CONFIG', 'sam2_hiera_l.yaml')
 MODEL_CHECKPOINT = os.getenv('MODEL_CHECKPOINT', 'sam2_hiera_large.pt')
 MAX_FRAMES_TO_TRACK = int(os.getenv('MAX_FRAMES_TO_TRACK', 10))
@@ -35,7 +34,7 @@
 
 
 # build path to the model checkpoint
-sam2_checkpoint = str(pathlib.Path(__file__).parent / SEGMENT_ANYTHING_2_REPO_PATH / "checkpoints" / MODEL_CHECKPOINT)
+sam2_checkpoint = str(pathlib.Path(__file__).parent / "/sam2" / "checkpoints" / MODEL_CHECKPOINT)
 predictor = build_sam2_video_predictor(MODEL_CONFIG, sam2_checkpoint)