Add Docker Support for AmadeusGPT (SAM2) and SAM Image Segmentation

Dockerfile

@@ -0,0 +1,54 @@
+FROM pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtime
+ARG user_name
+ARG uid
+ARG gid
+ENV HOME=/root 
+WORKDIR ${HOME}
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update -y && apt-get install -y --no-install-recommends \
+    vim zsh tmux wget curl htop jupyter python3 python3-pip libgl1-mesa-glx git sudo ssh libglib2.0-0 \
+    tree \
+    && apt-get -y autoclean \
+    && apt-get -y autoremove \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* \
+    && apt-get clean \
+    && apt-get update 
+
+# Install conda packages as root
+SHELL ["/bin/bash", "-lc"]
+RUN conda install -y python=3.10 pytables=3.8.0 hdf5 jupyter
+
+# Install DeepLabCut and other pip packages as root
+RUN pip install --no-cache-dir --upgrade pip setuptools && \
+    pip install --no-cache-dir gpustat nvitop pytest PySide6==6.3.1 streamlit networkx isort black && \
+    pip install --no-cache-dir amadeusgpt && \
+    pip install --no-cache-dir git+https://github.com/DeepLabCut/DeepLabCut.git
+
+# Initialize conda for zsh shell
+RUN /opt/conda/bin/conda init zsh
+
+# Create a non-root user
+RUN mkdir /app /logs /data
+RUN groupadd -g ${gid} ${user_name} \
+    && useradd -m -u ${uid} -g ${gid} ${user_name} \
+    && chown -R ${uid}:${gid} /home
+
+# Switch to the new user and set home directory as working directory
+USER ${user_name}
+ENV HOME=/home/${user_name}
+WORKDIR ${HOME}
+
+# Install Oh My Zsh and plugins
+RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" --unattended \
+    && git clone https://github.com/zsh-users/zsh-autosuggestions ~/.oh-my-zsh/custom/plugins/zsh-autosuggestions \
+    && git clone https://github.com/zsh-users/zsh-syntax-highlighting ~/.oh-my-zsh/custom/plugins/zsh-syntax-highlighting \
+    && sed -i 's/^plugins=(.*)$/plugins=(git zsh-autosuggestions zsh-syntax-highlighting)/' ~/.zshrc
+
+RUN echo "export PATH=$PATH:/home/${user_name}/.local/bin" >> /home/${user_name}/.zshrc
+
+USER root
+WORKDIR /app
+
+CMD ["zsh"]
+SHELL ["/bin/zsh", "-c"]

Dockerfile.sam2

@@ -0,0 +1,57 @@
+FROM pytorch/pytorch:2.3.1-cuda11.8-cudnn8-runtime
+ARG user_name
+ARG uid
+ARG gid
+ENV HOME=/root 
+WORKDIR ${HOME}
+
+ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update -y && apt-get install -y --no-install-recommends \
+    vim zsh tmux wget htop jupyter python3 python3-pip libgl1-mesa-glx git sudo ssh libglib2.0-0 \
+    tree \
+    && apt-get -y autoclean \
+    && apt-get -y autoremove \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* \
+    && apt-get clean \
+    && apt-get update 
+
+# Install conda packages as root
+SHELL ["/bin/bash", "-lc"]
+RUN conda install -y python=3.10 pytables=3.8.0 hdf5 jupyter
+
+# Install DeepLabCut and other pip packages as root
+RUN pip install --no-cache-dir --upgrade pip setuptools && \
+    pip install --no-cache-dir gpustat nvitop pytest PySide6==6.3.1 streamlit networkx isort black && \
+    pip install --no-cache-dir amadeusgpt && \
+    pip install --no-cache-dir git+https://github.com/DeepLabCut/DeepLabCut.git
+
+# Initialize conda for zsh shell
+RUN /opt/conda/bin/conda init zsh
+
+# Create a non-root user
+RUN mkdir /app /logs /data
+RUN groupadd -g ${gid} ${user_name} \
+    && useradd -m -u ${uid} -g ${gid} ${user_name} \
+    && chown -R ${uid}:${gid} /home
+
+# Switch to the new user and set home directory as working directory
+USER ${user_name}
+ENV HOME=/home/${user_name}
+WORKDIR ${HOME}
+
+RUN git clone https://github.com/facebookresearch/sam2.git SAM2 && cd SAM2 && pip install -e . && cd checkpoints && sh download_ckpts.sh
+
+# Install Oh My Zsh and plugins
+# RUN sh -c "$(curl -fsSL https://raw.githubusercontent.com/ohmyzsh/ohmyzsh/master/tools/install.sh)" --unattended \
+#     && git clone https://github.com/zsh-users/zsh-autosuggestions ~/.oh-my-zsh/custom/plugins/zsh-autosuggestions \
+#     && git clone https://github.com/zsh-users/zsh-syntax-highlighting ~/.oh-my-zsh/custom/plugins/zsh-syntax-highlighting \
+#     && sed -i 's/^plugins=(.*)$/plugins=(git zsh-autosuggestions zsh-syntax-highlighting)/' ~/.zshrc
+# RUN echo "export PATH=$PATH:/home/${user_name}/.local/bin" >> /home/${user_name}/.zshrc
+
+USER root
+WORKDIR ${HOME}
+
+# CMD ["zsh"]
+# SHELL ["/bin/zsh", "-c"]
+CMD ["/bin/bash"]
+SHELL ["/bin/bash", "-c"]

Makefile

@@ -2,3 +2,36 @@ export streamlit_app=True
 app:
 
 	streamlit run amadeusgpt/app.py --server.fileWatcherType none --server.maxUploadSize 1000
+
+IMG_TAG := 0.2
+IMG_NAME := amadeusgpt
+DOCKERFILE := Dockerfile.sam2
+
+BUILD_ARGS := \
+        --build-arg uid=$(shell id -u) \
+        --build-arg gid=$(shell id -g) \
+        --build-arg user_name=$(shell id -un)
+build:
+	docker build $(BUILD_ARGS) \
+		 -t $(IMG_NAME):$(IMG_TAG) -f $(DOCKERFILE) .
+
+# [USER: ADJUST VOLUMES]
+# path to the local project
+HOST_SRC := /home/$(shell id -un)/AmadeusGPT
+# path to the project in the container
+DOCKER_SRC := /home/$(shell id -un)/AmadeusGPT
+# DOCKER_DATA := /mnt
+VOLUMES := \
+	--volume $(HOST_SRC):$(DOCKER_SRC) 
+
+CONTAINER_TAG :=v0.14
+CONTAINER_NAME := amadeusgpt_$(CONTAINER_TAG)
+
+run:
+	docker run --shm-size=60G --gpus all -it --name $(CONTAINER_NAME) \
+							$(VOLUMES) \
+							$(IMG_NAME):$(IMG_TAG)
+#  tail -f /dev/null
+
+exec:
+	docker exec -it $(CONTAINER_NAME) /bin/bash

sam2_image.py

@@ -0,0 +1,111 @@
+import torch
+from sam2.build_sam import build_sam2
+from sam2.sam2_image_predictor import SAM2ImagePredictor
+from PIL import Image
+import numpy as np
+from matplotlib import pyplot as plt
+import os
+
+def sam2_image_seg(image_path, prompt=None, output_path=None):
+    """
+    Segment an image using SAM2.
+
+    Args:
+        image_path (str): Path to the input image
+        prompt (dict, optional): Dictionary with prompts for the model.
+            Can contain 'point_coords', 'point_labels', and/or 'box'.
+            Example: {'point_coords': np.array([[x, y]]), 'point_labels': np.array([1])}
+        output_path (str, optional): Path to save visualization. If None, no visualization is saved.
+
+    Returns:
+        np.ndarray: Segmentation mask
+    """    
+    # Model paths
+    # checkpoint = "~/SAM2/checkpoints/sam2.1_hiera_small.pt"
+    checkpoint = os.path.expanduser("~/SAM2/checkpoints/sam2.1_hiera_small.pt")
+    model_cfg = "configs/sam2.1/sam2.1_hiera_s.yaml"
+
+    # Initialize predictor
+    predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint))
+
+    # Set device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    predictor.model = predictor.model.to(device)
+
+    # Load image
+    image = Image.open(image_path)
+    image = np.array(image.convert("RGB"))
+    predictor.set_image(image)
+
+    # Run prediction with appropriate precision
+    if device == "cuda":
+        with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+            if prompt is None:
+                masks, scores, logits = predictor.predict()
+            else:
+                masks, scores, logits = predictor.predict(
+                    point_coords=prompt.get('point_coords', None),
+                    point_labels=prompt.get('point_labels', None),
+                    box=prompt.get('box', None),
+                    multimask_output=True
+                )
+    else:
+        with torch.inference_mode():
+            if prompt is None:
+                masks, scores, logits = predictor.predict()
+            else:
+                masks, scores, logits = predictor.predict(
+                    point_coords=prompt.get('point_coords', None),
+                    point_labels=prompt.get('point_labels', None),
+                    box=prompt.get('box', None),
+                    multimask_output=True
+                )
+
+    # Save visualization if output_path is provided
+    if output_path is not None:
+        # Create output directory if it doesn't exist
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+
+        plt.figure(figsize=(10, 10))
+        plt.imshow(image)  # Show the original image
+        plt.imshow(masks[0], cmap="jet", alpha=0.5)  # Overlay the first mask with transparency
+        plt.axis("off")  # Remove axes for better visualization
+        plt.title("Image with Predicted Mask")
+        plt.savefig(output_path)
+        plt.close()
+
+    return masks
+
+# from SAM.sam2_image import sam2_image_seg
+# from sam2.build_sam import build_sam2
+
+if __name__ == "__main__":
+
+    # checkpoint = "./checkpoints/sam2.1_hiera_large.pt"
+    # model_cfg = "configs/sam2.1/sam2.1_hiera_l.yaml"
+    # checkpoint = "/home/ti_wang/AmadeusGPT/sam2/checkpoints/sam2.1_hiera_small.pt"
+    # model_cfg = "configs/sam2.1/sam2.1_hiera_s.yaml"
+
+    # predictor = SAM2ImagePredictor(build_sam2(model_cfg, checkpoint))
+
+    # with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16):
+    #     image_path = "./notebooks/images/cars.jpg"
+
+    #     image = Image.open(image_path)
+    #     image = np.array(image.convert("RGB"))
+    #     # image = np.array(image)
+
+    #     predictor.set_image(image)
+    #     masks, _, _ = predictor.predict()
+
+    #     # Plot the original image and overlay the mask
+    #     plt.figure(figsize=(10, 10))
+    #     plt.imshow(image)  # Show the original image
+    #     plt.imshow(masks[0], cmap="jet", alpha=0.5)  # Overlay the first mask with transparency
+    #     plt.axis("off")  # Remove axes for better visualization
+    #     plt.title("Image with Predicted Mask")
+    #     plt.savefig("./test_images/mask_overlay.png")  # Save the figure
+
+    image_path = "./SAM/notebooks/images/truck.jpg"
+    output_path = "./ti_test/mask_overlay_4.png"
+    sam2_image_seg(image_path, output_path=output_path)