diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000..988a428 --- /dev/null +++ b/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"] \ No newline at end of file diff --git a/Dockerfile.sam2 b/Dockerfile.sam2 new file mode 100755 index 0000000..6dbc5b7 --- /dev/null +++ b/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"] \ No newline at end of file diff --git a/Makefile b/Makefile old mode 100644 new mode 100755 index 56f5a7b..a0e3637 --- a/Makefile +++ b/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 diff --git a/sam2_image.py b/sam2_image.py new file mode 100755 index 0000000..404515d --- /dev/null +++ b/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) \ No newline at end of file