init

Dockerfile

@@ -0,0 +1,34 @@
+FROM python:3.10
+
+WORKDIR /code
+
+COPY ./requirements.txt /code/requirements.txt
+
+RUN apt-get update && apt-get install -y libgl1-mesa-glx
+RUN pip install --no-cache-dir torch
+RUN pip install --no-cache-dir torchvision
+RUN pip install --no-cache-dir git+https://github.com/luca-medeiros/lang-segment-anything.git
+RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+# Switch to the "user" user
+USER user
+# Set home to the user's home directory
+ENV HOME=/home/user \
+	PATH=/home/user/.local/bin:$PATH \
+    PYTHONPATH=$HOME/app \
+	PYTHONUNBUFFERED=1 \
+	GRADIO_ALLOW_FLAGGING=never \
+	GRADIO_NUM_PORTS=1 \
+	GRADIO_SERVER_NAME=0.0.0.0 \
+	GRADIO_THEME=huggingface \
+	SYSTEM=spaces
+ 
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+CMD ["python", "app.py"]

app.py

@@ -0,0 +1,63 @@
+# IMPORTS
+import torch
+import numpy as np
+from PIL import Image
+from lang_sam import LangSAM
+import gradio as gr
+
+
+def run_lang_sam(input_image, text_prompt, model):
+    image = input_image.convert("RGB").resize((512, 512))
+    masks, _, _, _ = model.predict(
+        image,
+        text_prompt
+    )
+    masks_int = masks.to(torch.uint8)
+    masks_max, _ = masks_int.max(dim=0, keepdim=True)
+    unified_mask = masks_max.squeeze(0).to(torch.bool)
+    return Image.fromarray(
+        (unified_mask[..., None].numpy() * np.array(image)).astype(np.uint8)
+    )
+
+
+def setup_gradio_interface(model):
+    block = gr.Blocks()
+
+    with block:
+        gr.Markdown("<h1><center>Lang SAM<h1><center>")
+
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(type="pil", label="Input Image")
+                text_prompt = gr.Textbox(label="Enter what you want to segment")
+                run_button = gr.Button(value="Run")
+
+            with gr.Column():
+                output_mask = gr.Image(type="numpy", label="Segmentation Mask")
+
+        run_button.click(
+            fn=lambda image, prompt: run_lang_sam(
+                image, prompt, model,
+            ),
+            inputs=[input_image, text_prompt],
+            outputs=[output_mask],
+        )
+
+        gr.Examples(
+            examples=[["bw-image.jpeg", "road"]],
+            inputs=[input_image, text_prompt],
+            outputs=[output_mask],
+            fn=lambda image, prompt: run_lang_sam(
+                image, prompt, model,
+            ),
+            cache_examples=True,
+            label="Try this example input!",
+        )
+
+    return block
+
+
+if __name__ == "__main__":
+    model = LangSAM()
+    gradio_interface = setup_gradio_interface(model)
+    gradio_interface.launch(share=False, show_api=False, show_error=True)

requirements.txt

@@ -0,0 +1 @@
+gradio==4.5.0