optimize-docker (#3)

- Add gitignore (c469934d2ab5ae27a4caeac02c0beb10a30a33ef)
- remove unused imports (346623e65d56fc39f4a1ab6bd708112b6eb8d3a6)
- optimise docker build (6eb6110ef4a3aff3660e6c2b5708cf0bc264c967)

.dockerignore

@@ -0,0 +1,10 @@
+env
+**/__pycache__
+
+.git/
+.github/
+.gitignore
+
+*.zip
+*.whl
+!gradio_image_prompter-0.1.0-py3-none-any.whl

.gitignore

@@ -0,0 +1,8 @@
+# python
+env/
+__pycache__
+.python-version
+
+
+# vim
+*.sw[op]

Dockerfile

@@ -1,45 +1,92 @@
+## Build time
 # Use the specified Python runtime as a parent image
-FROM docker.io/nvidia/cuda:12.1.0-cudnn8-devel-ubi8@sha256:f045009cab64c9fda6113b4473ac1c57dfcca65e18ce981bce63f3cddf7b807a
+FROM docker.io/nvidia/cuda:12.1.0-devel-ubuntu22.04@sha256:e3a8f7b933e77ecee74731198a2a5483e965b585cea2660675cf4bb152237e9b AS build
 
 # Set the working directory in the container
 WORKDIR /usr/src/app
+COPY packages.txt .
 
-# Install required packages
-RUN apt-get update && apt-get install -y \
-    gcc-11 \
-    build-essential \
-    ffmpeg \
-    libsm6 \
-    libxext6 \
-    curl \
-    git \
+ENV CUDA_HOME=/usr/local/cuda \
+    TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6+PTX" \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1
+
+
+# Delete nvidia apt list and Install required packages
+RUN DEBIAN_FRONTEND=noninteractive \
+    && rm -rf /etc/apt/sources.list.d/cuda.list /etc/apt/sources.list.d/nvidia-ml.list \
+    && apt-key del 7fa2af80 \
+    && apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub \
+    && apt-get -yq update \
+    && apt-get install --no-install-recommends -yq \
+        apt-transport-https \
+        ca-certificates \
+        $(cat packages.txt) \
+        python3 \
+        python3-dev \
+        python3-pip \
+        python3-venv \
+        ffmpeg \
+        libsm6 \
+        libxext6 \
     && apt-get clean \
     && rm -rf /var/lib/apt/lists/*
 
-# Set environment variable to use gcc-11
-ENV CC=/usr/bin/gcc-11
-
-# Copy the current directory contents into the container
-COPY . .
+# Setup virtual env
+RUN python3 -m venv /opt/venv
+ENV PATH="/opt/venv/bin:$PATH"
 
 # Install any needed packages specified in requirements.txt
-RUN pip install --no-cache-dir -r requirements.txt
+COPY requirements.txt .
+RUN --mount=type=cache,id=pip,target=/root/.cache \
+    pip install -r requirements.txt
 
-# Set the working directory for the GroundingDINO ops
+# Copy grounding dino ops
 WORKDIR /usr/src/app/models/GroundingDINO/ops
+COPY models/GroundingDINO/ops .
 
 # Run the setup script and the test script
-RUN python setup.py build install
-RUN python test.py # This should result in 6 lines of * True
+RUN CC=/usr/bin/gcc-11 python3 setup.py build && \
+    pip install .
+
+## Runtime
+# Use the specified Python runtime as a parent image
+FROM ubuntu:22.04
 
-# Install Gradio
-RUN pip install gradio
+ENV CUDA_HOME=/usr/local/cuda \
+    TORCH_CUDA_ARCH_LIST="6.0 6.1 7.0 7.5 8.0 8.6+PTX" \
+    PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PATH="/opt/venv/bin:/home/user/.local/bin:$PATH" \
+    HOME=/home/user
 
-# Change back to the original working directory
-WORKDIR /usr/src/app
+RUN DEBIAN_FRONTEND=noninteractive apt-get -yq update && apt-get install --no-install-recommends -yq \
+    python3 \
+    python3-dev \
+    python3-pip \
+    ffmpeg \
+    libsm6 \
+    libxext6 \
+    && apt-get clean \
+    && rm -rf /var/lib/apt/lists/*
 
+WORKDIR /app
+RUN useradd -m -u 1000 user && chown -R user /app
+
+COPY --from=build --chown=user /opt/venv /opt/venv
+
+COPY --chown=user checkpoints checkpoints
+COPY --chown=user checkpoint_best_regular.pth .
+COPY --chown=user *.jpg *.JPG ./
+COPY --chown=user datasets datasets
+COPY --chown=user groundingdino groundingdino
+COPY --chown=user models models
+COPY --chown=user util util
+COPY --chown=user app.py cfg_app.py ./
+
+USER user
 # Expose the port Gradio will run on
 EXPOSE 7860
-
+ENV GRADIO_SERVER_NAME="0.0.0.0"
 # Default command to run the Gradio app
-CMD ["python", "app.py"]
+CMD ["/opt/venv/bin/python3", "app.py"]

README.md

@@ -1,7 +1,22 @@
 ---
 title: CountGD_Multi-Modal_Open-World_Counting
-app_file: app.py
-sdk: gradio
-sdk_version: 4.44.1
+sdk: docker
+app_port: 7860
 ---
-nohup python -u app.py &
+# CountGD: Multi Modal Open World Counting Model
+
+To Run Locally, the best method is to use docker.
+
+Make sure you have installed docker, nvidia-driver and nvidia container toolkit for the your platform.
+
+Then, you can run the app locally with the following command
+
+```bash
+docker run -it \
+    --name countgd \
+    -p 7860:7860 \
+    --platform=linux/amd64 \
+    --gpus all \
+	registry.hf.space/nikigoli-countgd:latest \
+    python app.py
+```

app.py

@@ -1,16 +1,10 @@
 import spaces
 import gradio as gr
-import copy
 import random
 import torch
-import PIL
-from PIL import Image, ImageDraw, ImageFont
-import torchvision.transforms.functional as F
+from PIL import Image
 import numpy as np
 import argparse
-import json
-import plotly.express as px
-import pandas as pd
 from util.slconfig import SLConfig, DictAction
 from util.misc import nested_tensor_from_tensor_list
 import datasets.transforms as T
@@ -30,53 +24,9 @@ cwd = os.getcwd()
 import warnings
 warnings.filterwarnings("ignore")
 
-# Installing dependencies not in requirements.txt
-subprocess.run(
-    shlex.split(
-        "pip install gradio_image_prompter-0.1.0-py3-none-any.whl"
-    )
-)
 from gradio_image_prompter import ImagePrompter
-"""
-subprocess.run(
-    shlex.split(
-        "pip install MultiScaleDeformableAttention-1.0-cp310-cp310-linux_x86_64.whl"
-    )
-)
-"""
-"""
-subprocess.run(
-    shlex.split(
-        "python test.py"
-    )
-)
-"""
-#with open('./switch_cuda.sh', 'rb') as file:
-#    script = file.read()
-#call(script, shell=True)
-
-with open('./build_ops.sh', 'rb') as file:
-    script = file.read()
-call(script, shell=True)
-
-def find_cuda():
-    # Check if CUDA_HOME or CUDA_PATH environment variables are set
-    cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH')
-
-    if cuda_home and os.path.exists(cuda_home):
-        return cuda_home
-
-    # Search for the nvcc executable in the system's PATH
-    nvcc_path = shutil.which('nvcc')
-
-    if nvcc_path:
-        # Remove the 'bin/nvcc' part to get the CUDA installation path
-        cuda_path = os.path.dirname(os.path.dirname(nvcc_path))
-        return cuda_path
 
-    return None
 
-cuda_path = find_cuda()
 
 class AppSteps(Enum):
     JUST_TEXT = 1
@@ -195,13 +145,6 @@ def build_model_and_transforms(args):
 
     return model, data_transform
 
-
-parser = argparse.ArgumentParser("Counting Application", parents=[get_args_parser()])
-args = parser.parse_args()
-device = get_device()
-model, transform = build_model_and_transforms(args)
-model = model.to(device)
-
 examples = [
     ["strawberry.jpg", "strawberry", {"image": "strawberry.jpg"}],
     ["strawberry.jpg", "blueberry", {"image": "strawberry.jpg"}],
@@ -254,242 +197,250 @@ def get_ind_to_filter(text, word_ids, keywords):
 
     return inds_to_filter
 
-@spaces.GPU(duration=120)
-def count(image, text, prompts, state, device):
-
-    keywords = "" # do not handle this for now
-    
-    # Handle no prompt case.
-    if prompts is None:
-        prompts = {"image": image, "points": []}
-    input_image, _ = transform(image, {"exemplars": torch.tensor([])})
-    input_image = input_image.unsqueeze(0).to(device)
-    exemplars = get_box_inputs(prompts["points"])
-    
-    input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
-    input_image_exemplars = input_image_exemplars.unsqueeze(0).to(device)
-    exemplars = [exemplars["exemplars"].to(device)]
-
-    with torch.no_grad():
-        model_output = model(
-                nested_tensor_from_tensor_list(input_image),
-                nested_tensor_from_tensor_list(input_image_exemplars),
-                exemplars,
-                [torch.tensor([0]).to(device) for _ in range(len(input_image))],
-                captions=[text + " ."] * len(input_image),
-            )
-    
-    ind_to_filter = get_ind_to_filter(text, model_output["token"][0].word_ids, keywords)
-    logits = model_output["pred_logits"].sigmoid()[0][:, ind_to_filter]
-    boxes = model_output["pred_boxes"][0]
-    if len(keywords.strip()) > 0:
-        box_mask = (logits > CONF_THRESH).sum(dim=-1) == len(ind_to_filter)
-    else:
-        box_mask = logits.max(dim=-1).values > CONF_THRESH
-    logits = logits[box_mask, :].cpu().numpy()
-    boxes = boxes[box_mask, :].cpu().numpy()
-    
-    # Plot results.
-    (w, h) = image.size
-    det_map = np.zeros((h, w))
-    det_map[(h * boxes[:, 1]).astype(int), (w * boxes[:, 0]).astype(int)] = 1
-    det_map = ndimage.gaussian_filter(
-        det_map, sigma=(w // 200, w // 200), order=0
-    )
-    plt.imshow(image)
-    plt.imshow(det_map[None, :].transpose(1, 2, 0), 'jet', interpolation='none', alpha=0.7)
-    plt.axis('off')
-    img_buf = io.BytesIO()
-    plt.savefig(img_buf, format='png', bbox_inches='tight')
-    plt.close()
-
-    output_img = Image.open(img_buf)
-
-    if AppSteps.TEXT_AND_EXEMPLARS not in state:
-        exemplar_image = ImagePrompter(type='pil', label='Visual Exemplar Image', value=prompts, interactive=True, visible=True)
-        new_submit_btn = gr.Button("Count", variant="primary", interactive=False)
-        state = [AppSteps.JUST_TEXT, AppSteps.TEXT_AND_EXEMPLARS]
-        main_instructions_comp = gr.Markdown(visible=False)
-        step_3 = gr.Tab(visible=False)
-    elif AppSteps.FULL_APP not in state:
-        exemplar_image = ImagePrompter(type='pil', label='Visual Exemplar Image', value=prompts, interactive=True, visible=True)
-        new_submit_btn = submit_btn
-        state = [AppSteps.JUST_TEXT, AppSteps.TEXT_AND_EXEMPLARS, AppSteps.FULL_APP]
-        main_instructions_comp = gr.Markdown(visible=True)
-        step_3 = gr.Tab(visible=True)
-    else:
-        exemplar_image = ImagePrompter(type='pil', label='Visual Exemplar Image', value=prompts, interactive=True, visible=True)
-        new_submit_btn = submit_btn
-        main_instructions_comp = gr.Markdown(visible=True)
-        step_3 = gr.Tab(visible=True)
-
-    out_label = "Detected instances predicted with"
-    if len(text.strip()) > 0:
-        out_label += " text"
-        if exemplars[0].size()[0] == 1:
-            out_label += " and " + str(exemplars[0].size()[0]) + " visual exemplar."        
-        elif exemplars[0].size()[0] > 1:
-            out_label += " and " + str(exemplars[0].size()[0]) + " visual exemplars."
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser("Counting Application", parents=[get_args_parser()])
+    args = parser.parse_args()
+    device = get_device()
+    model, transform = build_model_and_transforms(args)
+    model = model.to(device)
+
+    @spaces.GPU(duration=120)
+    def count(image, text, prompts, state, device):
+
+        keywords = "" # do not handle this for now
+
+        # Handle no prompt case.
+        if prompts is None:
+            prompts = {"image": image, "points": []}
+        input_image, _ = transform(image, {"exemplars": torch.tensor([])})
+        input_image = input_image.unsqueeze(0).to(device)
+        exemplars = get_box_inputs(prompts["points"])
+
+        input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
+        input_image_exemplars = input_image_exemplars.unsqueeze(0).to(device)
+        exemplars = [exemplars["exemplars"].to(device)]
+
+        with torch.no_grad():
+            model_output = model(
+                    nested_tensor_from_tensor_list(input_image),
+                    nested_tensor_from_tensor_list(input_image_exemplars),
+                    exemplars,
+                    [torch.tensor([0]).to(device) for _ in range(len(input_image))],
+                    captions=[text + " ."] * len(input_image),
+                )
+
+        ind_to_filter = get_ind_to_filter(text, model_output["token"][0].word_ids, keywords)
+        logits = model_output["pred_logits"].sigmoid()[0][:, ind_to_filter]
+        boxes = model_output["pred_boxes"][0]
+        if len(keywords.strip()) > 0:
+            box_mask = (logits > CONF_THRESH).sum(dim=-1) == len(ind_to_filter)
         else:
-            out_label += "."
-    elif exemplars[0].size()[0] > 0:
-        if exemplars[0].size()[0] == 1:
-            out_label += " " + str(exemplars[0].size()[0]) + " visual exemplar."
+            box_mask = logits.max(dim=-1).values > CONF_THRESH
+        logits = logits[box_mask, :].cpu().numpy()
+        boxes = boxes[box_mask, :].cpu().numpy()
+
+        # Plot results.
+        (w, h) = image.size
+        det_map = np.zeros((h, w))
+        det_map[(h * boxes[:, 1]).astype(int), (w * boxes[:, 0]).astype(int)] = 1
+        det_map = ndimage.gaussian_filter(
+            det_map, sigma=(w // 200, w // 200), order=0
+        )
+        plt.imshow(image)
+        plt.imshow(det_map[None, :].transpose(1, 2, 0), 'jet', interpolation='none', alpha=0.7)
+        plt.axis('off')
+        img_buf = io.BytesIO()
+        plt.savefig(img_buf, format='png', bbox_inches='tight')
+        plt.close()
+
+        output_img = Image.open(img_buf)
+
+        if AppSteps.TEXT_AND_EXEMPLARS not in state:
+            exemplar_image = ImagePrompter(type='pil', label='Visual Exemplar Image', value=prompts, interactive=True, visible=True)
+            new_submit_btn = gr.Button("Count", variant="primary", interactive=False)
+            state = [AppSteps.JUST_TEXT, AppSteps.TEXT_AND_EXEMPLARS]
+            main_instructions_comp = gr.Markdown(visible=False)
+            step_3 = gr.Tab(visible=False)
+        elif AppSteps.FULL_APP not in state:
+            exemplar_image = ImagePrompter(type='pil', label='Visual Exemplar Image', value=prompts, interactive=True, visible=True)
+            new_submit_btn = submit_btn
+            state = [AppSteps.JUST_TEXT, AppSteps.TEXT_AND_EXEMPLARS, AppSteps.FULL_APP]
+            main_instructions_comp = gr.Markdown(visible=True)
+            step_3 = gr.Tab(visible=True)
         else:
-            out_label += " " + str(exemplars[0].size()[0]) + " visual exemplars."
-    else:
-        out_label = "Nothing specified to detect."
- 
-    return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]), new_submit_btn, gr.Tab(visible=True), step_3, state)
-
-@spaces.GPU
-def count_main(image, text, prompts, device):
-    keywords = "" # do not handle this for now
-    # Handle no prompt case.
-    if prompts is None:
-        prompts = {"image": image, "points": []}
-    input_image, _ = transform(image, {"exemplars": torch.tensor([])})
-    input_image = input_image.unsqueeze(0).to(device)
-    exemplars = get_box_inputs(prompts["points"])
-    
-    input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
-    input_image_exemplars = input_image_exemplars.unsqueeze(0).to(device)
-    exemplars = [exemplars["exemplars"].to(device)]
-    
-    with torch.no_grad():
-        model_output = model(
-                nested_tensor_from_tensor_list(input_image),
-                nested_tensor_from_tensor_list(input_image_exemplars),
-                exemplars,
-                [torch.tensor([0]).to(device) for _ in range(len(input_image))],
-                captions=[text + " ."] * len(input_image),
-            )
-    
-    ind_to_filter = get_ind_to_filter(text, model_output["token"][0].word_ids, keywords)
-    logits = model_output["pred_logits"].sigmoid()[0][:, ind_to_filter]
-    boxes = model_output["pred_boxes"][0]
-    if len(keywords.strip()) > 0:
-        box_mask = (logits > CONF_THRESH).sum(dim=-1) == len(ind_to_filter)
-    else:
-        box_mask = logits.max(dim=-1).values > CONF_THRESH
-    logits = logits[box_mask, :].cpu().numpy()
-    boxes = boxes[box_mask, :].cpu().numpy()
-    
-    # Plot results.
-    (w, h) = image.size
-    det_map = np.zeros((h, w))
-    det_map[(h * boxes[:, 1]).astype(int), (w * boxes[:, 0]).astype(int)] = 1
-    det_map = ndimage.gaussian_filter(
-        det_map, sigma=(w // 200, w // 200), order=0
-    )
-    plt.imshow(image)
-    plt.imshow(det_map[None, :].transpose(1, 2, 0), 'jet', interpolation='none', alpha=0.7)
-    plt.axis('off')
-    img_buf = io.BytesIO()
-    plt.savefig(img_buf, format='png', bbox_inches='tight')
-    plt.close()
-
-    output_img = Image.open(img_buf)
-
-    out_label = "Detected instances predicted with"
-    if len(text.strip()) > 0:
-        out_label += " text"
-        if exemplars[0].size()[0] == 1:
-            out_label += " and " + str(exemplars[0].size()[0]) + " visual exemplar."        
-        elif exemplars[0].size()[0] > 1:
-            out_label += " and " + str(exemplars[0].size()[0]) + " visual exemplars."
+            exemplar_image = ImagePrompter(type='pil', label='Visual Exemplar Image', value=prompts, interactive=True, visible=True)
+            new_submit_btn = submit_btn
+            main_instructions_comp = gr.Markdown(visible=True)
+            step_3 = gr.Tab(visible=True)
+
+        out_label = "Detected instances predicted with"
+        if len(text.strip()) > 0:
+            out_label += " text"
+            if exemplars[0].size()[0] == 1:
+                out_label += " and " + str(exemplars[0].size()[0]) + " visual exemplar."
+            elif exemplars[0].size()[0] > 1:
+                out_label += " and " + str(exemplars[0].size()[0]) + " visual exemplars."
+            else:
+                out_label += "."
+        elif exemplars[0].size()[0] > 0:
+            if exemplars[0].size()[0] == 1:
+                out_label += " " + str(exemplars[0].size()[0]) + " visual exemplar."
+            else:
+                out_label += " " + str(exemplars[0].size()[0]) + " visual exemplars."
         else:
-            out_label += "."
-    elif exemplars[0].size()[0] > 0:
-        if exemplars[0].size()[0] == 1:
-            out_label += " " + str(exemplars[0].size()[0]) + " visual exemplar."
+            out_label = "Nothing specified to detect."
+
+        return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]), new_submit_btn, gr.Tab(visible=True), step_3, state)
+
+    @spaces.GPU
+    def count_main(image, text, prompts, device):
+        keywords = "" # do not handle this for now
+        # Handle no prompt case.
+        if prompts is None:
+            prompts = {"image": image, "points": []}
+        input_image, _ = transform(image, {"exemplars": torch.tensor([])})
+        input_image = input_image.unsqueeze(0).to(device)
+        exemplars = get_box_inputs(prompts["points"])
+
+        input_image_exemplars, exemplars = transform(prompts["image"], {"exemplars": torch.tensor(exemplars)})
+        input_image_exemplars = input_image_exemplars.unsqueeze(0).to(device)
+        exemplars = [exemplars["exemplars"].to(device)]
+
+        with torch.no_grad():
+            model_output = model(
+                    nested_tensor_from_tensor_list(input_image),
+                    nested_tensor_from_tensor_list(input_image_exemplars),
+                    exemplars,
+                    [torch.tensor([0]).to(device) for _ in range(len(input_image))],
+                    captions=[text + " ."] * len(input_image),
+                )
+
+        ind_to_filter = get_ind_to_filter(text, model_output["token"][0].word_ids, keywords)
+        logits = model_output["pred_logits"].sigmoid()[0][:, ind_to_filter]
+        boxes = model_output["pred_boxes"][0]
+        if len(keywords.strip()) > 0:
+            box_mask = (logits > CONF_THRESH).sum(dim=-1) == len(ind_to_filter)
         else:
-            out_label += " " + str(exemplars[0].size()[0]) + " visual exemplars."
-    else:
-        out_label = "Nothing specified to detect."
-    
-    return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]))
-
-def remove_label(image):
-    return gr.Image(show_label=False)
-
-def check_submit_btn(exemplar_image_prompts, state):
-    if AppSteps.TEXT_AND_EXEMPLARS not in state or len(state) == 3:
-        return gr.Button("Count", variant="primary", interactive=True)
-    elif exemplar_image_prompts is None:
-        return gr.Button("Count", variant="primary", interactive=False)
-    elif len(get_box_inputs(exemplar_image_prompts["points"])) > 0:
-        return gr.Button("Count", variant="primary", interactive=True)
-    else:
-        return gr.Button("Count", variant="primary", interactive=False)
-
-exemplar_img_drawing_instructions_part_1 = '<p><strong>Congrats, you have counted the strawberries!</strong> You can also draw a box around the object you want to count. <strong>Click and drag the mouse on the image below to draw a box around one of the strawberries.</strong> You can click the back button in the top right of the image to delete the box and try again.<img src="file/button-legend.jpg" width="750"></p>'
-exemplar_img_drawing_instructions_part_2 = '<p>The boxes you draw are called \"visual exemplars,\" image examples of what you want the model to count. You can add more boxes around more examples of strawberries in the image above to increase the accuracy of the predicted count. You can also use strawberries from a different image to specify the object to count by uploading or pasting a new image above and drawing boxes around strawberries in it.</p>'
-instructions_main = """
-# How to Use the App
-As shown earlier, there are 3 ways to specify the object to count: (1) with text only, (2) with text and any number of boxes (i.e., "visual exemplars") around example objects, and (3) with visual exemplars only. What is being used is indicated in the top left of the output image. How to try each case is detailed below.
-<ol>
-  <li><strong>Text Only: </strong> Only provide text describing the object to count in the textbox titled "What would you like to count?" Delete all boxes drawn on the visual exemplar image.</li>
-  <li><strong>Text + Visual Exemplars: </strong> Provide text describing the object to count in the textbox titled "What would you like to count?" and draw at least one box around an example object in the visual exemplar image.</li>
-  <li><strong>Visual Exemplars Only: </strong> Remove all text in the textbox titled "What would you like to count?" and draw at least one box around an example object in the visual exemplar image.</li>
-</ol>
-## Click on the "App" tab at the top of the screen to exit the tutorial and start using the main app!
-"""
-
-with gr.Blocks(title="CountGD: Multi-Modal Open-World Counting", theme="soft", head="""<meta name="viewport" content="width=device-width, initial-scale=1, user-scalable=1">""") as demo:
-    state = gr.State(value=[AppSteps.JUST_TEXT])
-    device = gr.State(device)
-    with gr.Tab("Tutorial"):
-        with gr.Row():
-            with gr.Column():
-                with gr.Tab("Step 3", visible=False) as step_3:
-                    main_instructions = gr.Markdown(instructions_main)
-                with gr.Tab("Step 2", visible=False) as step_2:
-                    gr.Markdown(exemplar_img_drawing_instructions_part_1)
-                    exemplar_image = ImagePrompter(type='pil', label='Visual Exemplar Image', show_label=True, value={"image": "strawberry.jpg", "points": []}, interactive=True)
-                    with gr.Accordion("Open for Further Information", open=False):
-                        gr.Markdown(exemplar_img_drawing_instructions_part_2)
-                with gr.Tab("Step 1", visible=True) as step_1:
-                    input_image = gr.Image(type='pil', label='Input Image', show_label='True', value="strawberry.jpg", interactive=False, width="30vw")
-                    gr.Markdown('# Click "Count" to count the strawberries.')
-                
-            with gr.Column():
-                with gr.Tab("Output Image"):
-                    detected_instances = gr.Image(label="Detected Instances", show_label='True', interactive=False, visible=True, width="40vw")
-                    
-        with gr.Row():
-            input_text = gr.Textbox(label="What would you like to count?", value="strawberry", interactive=True)
-            pred_count = gr.Number(label="Predicted Count", visible=False)
-        submit_btn = gr.Button("Count", variant="primary", interactive=True)
-                
-        submit_btn.click(fn=remove_label, inputs=[detected_instances], outputs=[detected_instances]).then(fn=count, inputs=[input_image, input_text, exemplar_image, state, device], outputs=[detected_instances, pred_count, submit_btn, step_2, step_3, state])
-        exemplar_image.change(check_submit_btn, inputs=[exemplar_image, state], outputs=[submit_btn])
-    with gr.Tab("App", visible=True) as main_app:
-        
-        gr.Markdown(
-              """
-              # <center>CountGD: Multi-Modal Open-World Counting
-              <center><h3>Count objects with text, visual exemplars, or both together.</h3>
-              <h3>Scroll down to try more examples</h3>
-              <h3><a href='https://arxiv.org/abs/2407.04619' target='_blank' rel='noopener'>[paper]</a>
-                <a href='https://github.com/niki-amini-naieni/CountGD/' target='_blank' rel='noopener'>[code]</a></h3>
-              Limitation: this app does not support fine-grained counting based on attributes or visual grounding inputs yet. Note: if the exemplar and text conflict each other, both will be counted.</center>
-              """
-            )
-            
-        with gr.Row():
-            with gr.Column():
-              input_image_main = gr.Image(type='pil', label='Input Image', show_label='True', value="strawberry.jpg", interactive=True)
-              input_text_main = gr.Textbox(label="What would you like to count?", placeholder="", value="strawberry")
-              exemplar_image_main = ImagePrompter(type='pil', label='Visual Exemplar Image', show_label=True, value={"image": "strawberry.jpg", "points": []}, interactive=True)
-            with gr.Column():
-              detected_instances_main = gr.Image(label="Detected Instances", show_label='True', interactive=False)
-              pred_count_main = gr.Number(label="Predicted Count")
-              submit_btn_main = gr.Button("Count", variant="primary")
-              clear_btn_main = gr.ClearButton(variant="secondary")
-        gr.Examples(label="Examples: click on a row to load the example. Add visual exemplars by drawing boxes on the loaded \"Visual Exemplar Image.\"", examples=examples, inputs=[input_image_main, input_text_main, exemplar_image_main])
-        submit_btn_main.click(fn=remove_label, inputs=[detected_instances_main], outputs=[detected_instances_main]).then(fn=count_main, inputs=[input_image_main, input_text_main, exemplar_image_main, device], outputs=[detected_instances_main, pred_count_main])
-        clear_btn_main.add([input_image_main, input_text_main, exemplar_image_main, detected_instances_main, pred_count_main])
-        
-
-demo.queue().launch(allowed_paths=['back-icon.jpg', 'paste-icon.jpg', 'upload-icon.jpg', 'button-legend.jpg'])
+            box_mask = logits.max(dim=-1).values > CONF_THRESH
+        logits = logits[box_mask, :].cpu().numpy()
+        boxes = boxes[box_mask, :].cpu().numpy()
+
+        # Plot results.
+        (w, h) = image.size
+        det_map = np.zeros((h, w))
+        det_map[(h * boxes[:, 1]).astype(int), (w * boxes[:, 0]).astype(int)] = 1
+        det_map = ndimage.gaussian_filter(
+            det_map, sigma=(w // 200, w // 200), order=0
+        )
+        plt.imshow(image)
+        plt.imshow(det_map[None, :].transpose(1, 2, 0), 'jet', interpolation='none', alpha=0.7)
+        plt.axis('off')
+        img_buf = io.BytesIO()
+        plt.savefig(img_buf, format='png', bbox_inches='tight')
+        plt.close()
+
+        output_img = Image.open(img_buf)
+
+        out_label = "Detected instances predicted with"
+        if len(text.strip()) > 0:
+            out_label += " text"
+            if exemplars[0].size()[0] == 1:
+                out_label += " and " + str(exemplars[0].size()[0]) + " visual exemplar."
+            elif exemplars[0].size()[0] > 1:
+                out_label += " and " + str(exemplars[0].size()[0]) + " visual exemplars."
+            else:
+                out_label += "."
+        elif exemplars[0].size()[0] > 0:
+            if exemplars[0].size()[0] == 1:
+                out_label += " " + str(exemplars[0].size()[0]) + " visual exemplar."
+            else:
+                out_label += " " + str(exemplars[0].size()[0]) + " visual exemplars."
+        else:
+            out_label = "Nothing specified to detect."
+
+        return (gr.Image(output_img, visible=True, label=out_label, show_label=True), gr.Number(label="Predicted Count", visible=True, value=boxes.shape[0]))
+
+    def remove_label(image):
+        return gr.Image(show_label=False)
+
+    def check_submit_btn(exemplar_image_prompts, state):
+        if AppSteps.TEXT_AND_EXEMPLARS not in state or len(state) == 3:
+            return gr.Button("Count", variant="primary", interactive=True)
+        elif exemplar_image_prompts is None:
+            return gr.Button("Count", variant="primary", interactive=False)
+        elif len(get_box_inputs(exemplar_image_prompts["points"])) > 0:
+            return gr.Button("Count", variant="primary", interactive=True)
+        else:
+            return gr.Button("Count", variant="primary", interactive=False)
+
+    exemplar_img_drawing_instructions_part_1 = '<p><strong>Congrats, you have counted the strawberries!</strong> You can also draw a box around the object you want to count. <strong>Click and drag the mouse on the image below to draw a box around one of the strawberries.</strong> You can click the back button in the top right of the image to delete the box and try again.<img src="file/button-legend.jpg" width="750"></p>'
+    exemplar_img_drawing_instructions_part_2 = '<p>The boxes you draw are called \"visual exemplars,\" image examples of what you want the model to count. You can add more boxes around more examples of strawberries in the image above to increase the accuracy of the predicted count. You can also use strawberries from a different image to specify the object to count by uploading or pasting a new image above and drawing boxes around strawberries in it.</p>'
+    instructions_main = """
+    # How to Use the App
+    As shown earlier, there are 3 ways to specify the object to count: (1) with text only, (2) with text and any number of boxes (i.e., "visual exemplars") around example objects, and (3) with visual exemplars only. What is being used is indicated in the top left of the output image. How to try each case is detailed below.
+    <ol>
+    <li><strong>Text Only: </strong> Only provide text describing the object to count in the textbox titled "What would you like to count?" Delete all boxes drawn on the visual exemplar image.</li>
+    <li><strong>Text + Visual Exemplars: </strong> Provide text describing the object to count in the textbox titled "What would you like to count?" and draw at least one box around an example object in the visual exemplar image.</li>
+    <li><strong>Visual Exemplars Only: </strong> Remove all text in the textbox titled "What would you like to count?" and draw at least one box around an example object in the visual exemplar image.</li>
+    </ol>
+    ## Click on the "App" tab at the top of the screen to exit the tutorial and start using the main app!
+    """
+
+    with gr.Blocks(title="CountGD: Multi-Modal Open-World Counting", theme="soft", head="""<meta name="viewport" content="width=device-width, initial-scale=1, user-scalable=1">""") as demo:
+        state = gr.State(value=[AppSteps.JUST_TEXT])
+        device = gr.State(device)
+        with gr.Tab("Tutorial"):
+            with gr.Row():
+                with gr.Column():
+                    with gr.Tab("Step 3", visible=False) as step_3:
+                        main_instructions = gr.Markdown(instructions_main)
+                    with gr.Tab("Step 2", visible=False) as step_2:
+                        gr.Markdown(exemplar_img_drawing_instructions_part_1)
+                        exemplar_image = ImagePrompter(type='pil', label='Visual Exemplar Image', show_label=True, value={"image": "strawberry.jpg", "points": []}, interactive=True)
+                        with gr.Accordion("Open for Further Information", open=False):
+                            gr.Markdown(exemplar_img_drawing_instructions_part_2)
+                    with gr.Tab("Step 1", visible=True) as step_1:
+                        input_image = gr.Image(type='pil', label='Input Image', show_label='True', value="strawberry.jpg", interactive=False, width="30vw")
+                        gr.Markdown('# Click "Count" to count the strawberries.')
+
+                with gr.Column():
+                    with gr.Tab("Output Image"):
+                        detected_instances = gr.Image(label="Detected Instances", show_label='True', interactive=False, visible=True, width="40vw")
+
+            with gr.Row():
+                input_text = gr.Textbox(label="What would you like to count?", value="strawberry", interactive=True)
+                pred_count = gr.Number(label="Predicted Count", visible=False)
+            submit_btn = gr.Button("Count", variant="primary", interactive=True)
+
+            submit_btn.click(fn=remove_label, inputs=[detected_instances], outputs=[detected_instances]).then(fn=count, inputs=[input_image, input_text, exemplar_image, state, device], outputs=[detected_instances, pred_count, submit_btn, step_2, step_3, state])
+            exemplar_image.change(check_submit_btn, inputs=[exemplar_image, state], outputs=[submit_btn])
+        with gr.Tab("App", visible=True) as main_app:
+
+            gr.Markdown(
+                """
+                # <center>CountGD: Multi-Modal Open-World Counting
+                <center><h3>Count objects with text, visual exemplars, or both together.</h3>
+                <h3>Scroll down to try more examples</h3>
+                <h3><a href='https://arxiv.org/abs/2407.04619' target='_blank' rel='noopener'>[paper]</a>
+                    <a href='https://github.com/niki-amini-naieni/CountGD/' target='_blank' rel='noopener'>[code]</a></h3>
+                Limitation: this app does not support fine-grained counting based on attributes or visual grounding inputs yet. Note: if the exemplar and text conflict each other, both will be counted.</center>
+                """
+                )
+
+            with gr.Row():
+                with gr.Column():
+                    input_image_main = gr.Image(type='pil', label='Input Image', show_label='True', value="strawberry.jpg", interactive=True)
+                    input_text_main = gr.Textbox(label="What would you like to count?", placeholder="", value="strawberry")
+                    exemplar_image_main = ImagePrompter(type='pil', label='Visual Exemplar Image', show_label=True, value={"image": "strawberry.jpg", "points": []}, interactive=True)
+                with gr.Column():
+                    detected_instances_main = gr.Image(label="Detected Instances", show_label='True', interactive=False)
+                    pred_count_main = gr.Number(label="Predicted Count")
+                    submit_btn_main = gr.Button("Count", variant="primary")
+                    clear_btn_main = gr.ClearButton(variant="secondary")
+            gr.Examples(label="Examples: click on a row to load the example. Add visual exemplars by drawing boxes on the loaded \"Visual Exemplar Image.\"", examples=examples, inputs=[input_image_main, input_text_main, exemplar_image_main])
+            submit_btn_main.click(fn=remove_label, inputs=[detected_instances_main], outputs=[detected_instances_main]).then(fn=count_main, inputs=[input_image_main, input_text_main, exemplar_image_main, device], outputs=[detected_instances_main, pred_count_main])
+            clear_btn_main.add([input_image_main, input_text_main, exemplar_image_main, detected_instances_main, pred_count_main])
+
+
+    demo.queue().launch(allowed_paths=['back-icon.jpg', 'paste-icon.jpg', 'upload-icon.jpg', 'button-legend.jpg'])

models/GroundingDINO/ops/setup.py

@@ -11,8 +11,6 @@ import glob
 
 import torch
 
-import spaces
-
 from torch.utils.cpp_extension import CUDA_HOME
 from torch.utils.cpp_extension import CppExtension
 from torch.utils.cpp_extension import CUDAExtension
@@ -22,7 +20,6 @@ from setuptools import setup
 
 requirements = ["torch", "torchvision"]
 
-#@spaces.GPU
 def get_extensions():
     this_dir = os.path.dirname(os.path.abspath(__file__))
     extensions_dir = os.path.join(this_dir, "src")
@@ -36,12 +33,9 @@ def get_extensions():
     extra_compile_args = {"cxx": []}
     define_macros = []
 
-
-
     print("inside get_extensions")
-    print(torch.cuda.is_available())
     print(CUDA_HOME)
-    if torch.cuda.is_available():
+    if CUDA_HOME is not None and (torch.cuda.is_available() or ("TORCH_CUDA_ARCH_LIST" in os.environ) or torch.cuda.get_arch_list()):
         extension = CUDAExtension
         sources += source_cuda
         define_macros += [("WITH_CUDA", None)]

models/GroundingDINO/ops/test.py

@@ -10,9 +10,8 @@ from __future__ import absolute_import
 from __future__ import print_function
 from __future__ import division
 
-import time
+
 import torch
-import torch.nn as nn
 from torch.autograd import gradcheck
 
 from functions.ms_deform_attn_func import MSDeformAttnFunction, ms_deform_attn_core_pytorch

packages.txt

@@ -1,2 +1,3 @@
 build-essential
 ninja-build
+gcc-11

requirements.txt

@@ -1,20 +1,18 @@
-cython
-submitit
 scipy
 termcolor
 addict
 yapf==0.40.1
 timm
-torch
-torchvision
-transformers
 numpy
 opencv-python
-supervision==0.6.0
 pycocotools
-pyyaml>3.10
 colorlog
-plotly-express
 setuptools
-wheel
-ushlex
+ushlex
+gradio>=4.0.0,<5
+gradio-image-prompter
+spaces
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch
+torchvision
+transformers