Add cp11 wheel for faster colab launch

app.py

@@ -176,27 +176,11 @@ def get_ind_to_filter(text, word_ids, keywords):
 
     return inds_to_filter
 
-def get_xy_from_boxes(image, boxes):
-    """
-    Get box centers and return in image coordinates
-    """
-    (w, h) = image.size
-    x = w * boxes[:, 0]
-    y = h * boxes[:, 1]
-
-    return x, y
-
 def generate_heatmap(image, boxes):
     # Plot results.
     (w, h) = image.size
     det_map = np.zeros((h, w))
-    x, y = get_xy_from_boxes(image, boxes)
-
-    # Box centers are floating point, convert to int and clip them at edge of box
-    x = np.clip(np.around(x).astype(int), 0, w - 1)
-    y = np.clip(np.around(y).astype(int), 0, h - 1)
-
-    det_map[y, x] = 1
+    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
     )
@@ -245,9 +229,13 @@ def preprocess(transform, image, input_prompts = None):
     return input_image, input_image_exemplar, exemplar
 
 def get_boxes_from_prediction(model_output, text, keywords = ""):
-    logits = model_output["pred_logits"].sigmoid()[0][:, :]
+    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]
-    box_mask = logits.max(dim=-1).values > CONF_THRESH
+    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
     boxes = boxes[box_mask, :].cpu().numpy()
     logits = logits[box_mask, :].cpu().numpy()
     return boxes, logits
@@ -426,5 +414,4 @@ if __name__ == '__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'])
-    #demo.queue().launch(share=True)
+    demo.queue().launch(allowed_paths=['back-icon.jpg', 'paste-icon.jpg', 'upload-icon.jpg', 'button-legend.jpg'])

notebooks/demo.ipynb

@@ -6,7 +6,7 @@
         "id": "yxig5CdZuHb9"
       },
       "source": [
-        "# CountGD - Multimodal open-world object counting\n",
+        "# CountGD - Multimodela open-world object counting\n",
         "\n"
       ]
     },
@@ -89,13 +89,7 @@
       "source": [
         "### Install Dependencies\n",
         "\n",
-        "The environment will be setup with the code, models and required dependencies.\n",
-        "\n",
-        "*Note for Colab users*\n",
-        "\n",
-        "To reduce the waiting time, you can use the pre-built wheel file available [here](https://drive.google.com/file/d/1Vl_6DAWfnVU7HFX5y_5TqqbkyTcjONbm/view?usp=sharing) - Visit the link and add it as a shortcut to your \"My Drive\" folder or edit the path accordingly below. (Line 28)\n",
-        "\n",
-        "Alternatively, if you are unable to use google drive, you can download the file to your machine & upload it to the colab runtime when you connect to it and update the path below to install it from there. (Line 28)"
+        "The environment will be setup with the code, models and required dependencies."
       ]
     },
     {
@@ -120,10 +114,8 @@
         "    git clone \"https://huggingface.co/spaces/nikigoli/countgd\" /content/countgd\n",
         "  fi\n",
         "  cd /content/countgd\n",
-        "\n",
-        "  # If you are testing out WIP items, uncomment the following and change the pr ref\n",
-        "  # git fetch origin refs/pr/10:refs/remotes/origin/pr/10\n",
-        "  # git checkout pr/10 && git pull\n",
+        "  git fetch origin refs/pr/5:refs/remotes/origin/pr/5\n",
+        "  git checkout pr/5\n",
         "else\n",
         "  # TODO check if cwd is the correct git repo\n",
         "  # If users use vscode, then we set the default start directory to root of the repo\n",
@@ -136,20 +128,11 @@
         "pip install --upgrade pip setuptools wheel\n",
         "pip install -r requirements.txt\n",
         "\n",
+        "# Compile modules\n",
+        "export CUDA_HOME=/usr/local/cuda/\n",
         "cd models/GroundingDINO/ops\n",
-        "if [ \"${RUNNING_IN_COLAB}\" == \"True\" ]; then\n",
-        "    export CUDA_HOME=/usr/local/cuda/\n",
-        "    if ! pip install \"/content/drive/MyDrive/MultiScaleDeformableAttention-1.0-cp311-cp311-linux_x86_64.whl\"\n",
-        "    then\n",
-        "        echo \"failed to install wheel, trying to build from source\";\n",
-        "        python3 setup.py build\n",
-        "        pip install .\n",
-        "    fi\n",
-        "else\n",
-        "    # We try to build the module as we dont know what environment we are running on\n",
-        "    python3 setup.py build\n",
-        "    pip install .\n",
-        "fi\n",
+        "python3 setup.py build\n",
+        "pip install .\n",
         "python3 test.py"
       ]
     },
@@ -206,7 +189,6 @@
         "    get_device,\n",
         "    get_args_parser,\n",
         "    generate_heatmap,\n",
-        "    get_xy_from_boxes,\n",
         "    predict,\n",
         ")\n",
         "args = get_args_parser().parse_args([])\n",
@@ -215,7 +197,7 @@
         "model = model.to(device)\n",
         "\n",
         "run = lambda image, text: predict(model, transform, image, text, None, device)\n",
-        "get_output = lambda image, boxes: (len(boxes), get_xy_from_boxes(image, boxes), generate_heatmap(image, boxes))\n"
+        "get_output = lambda image, boxes: (len(boxes), generate_heatmap(image, boxes))\n"
       ]
     },
     {
@@ -289,7 +271,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 15,
       "metadata": {
         "id": "rFXRk-_uuHb_"
       },
@@ -297,17 +279,6 @@
       "source": [
         "from tqdm import tqdm\n",
         "import os\n",
-        "import json\n",
-        "def convert_xy_to_json(xy: tuple):\n",
-        "    x, y = xy\n",
-        "    pts = []\n",
-        "    for _x, _y in zip(x.tolist(), y.tolist()):\n",
-        "        _x, _y = round(_x, 3), round(_y, 3)\n",
-        "        pts.append([_x, _y])\n",
-        "\n",
-        "    # List of [x, y] points\n",
-        "    return pts\n",
-        "\n",
         "def process_zipfile(input_zipfile: Path, text: str):\n",
         "    if not input_zipfile.exists() or not input_zipfile.is_file() or not os.access(input_zipfile, os.R_OK):\n",
         "        logger.error(f'Cannot open / read zipfile: {input_zipfile}. Please check if it exists')\n",
@@ -319,24 +290,14 @@
         "\n",
         "    output_zipfile = input_zipfile.parent / f'{input_zipfile.stem}_countgd.zip'\n",
         "    output_csvfile = input_zipfile.parent / f'{input_zipfile.stem}.csv'\n",
-        "    output_xyjson = input_zipfile.parent / f'{input_zipfile.stem}_xy.json'\n",
-        "\n",
-        "    xy_map = {}\n",
         "\n",
         "    logger.info(f'Writing outputs to {output_zipfile.name} and {output_csvfile.name} in {input_zipfile.parent} folder')\n",
         "    with zipfile_writer(output_zipfile) as add_to_zip, csvfile_writer(output_csvfile) as write_row:\n",
         "        for filename, im in tqdm(images_from_zipfile(input_zipfile)):\n",
-        "            try:\n",
-        "                boxes, _ = run(im, text)\n",
-        "                count, xy, heatmap  = get_output(im, boxes)\n",
-        "                logger.info(f'Count: {count} - {filename}')\n",
-        "                xy_map[filename] = convert_xy_to_json(xy)\n",
-        "                write_row({'filename': filename, 'count': count})\n",
-        "                add_to_zip(heatmap, filename)\n",
-        "            except Exception:\n",
-        "                logger.error(f'failed to process {filename}')\n",
-        "\n",
-        "    output_xyjson.write_text(json.dumps(xy_map))"
+        "            boxes, _ = run(im, text)\n",
+        "            count, heatmap  = get_output(im, boxes)\n",
+        "            write_row({'filename': filename, 'count': count})\n",
+        "            add_to_zip(heatmap, filename)"
       ]
     },
     {
@@ -349,8 +310,6 @@
         "\n",
         "Use the form on colab to set the parameters, providing the zipfile with input images and a promt text representing the object you want to count.\n",
         "\n",
-        "Use the fileupload button to add the zip file to the `countgd` directory or change the path below accordingly.\n",
-        "\n",
         "If you are not running on colab, change the values in the next cell\n",
         "\n",
         "Make sure to run the cell once you change the value."

requirements.txt

@@ -17,5 +17,4 @@ tqdm
 --extra-index-url https://download.pytorch.org/whl/cu121
 torch<2.6
 torchvision
-transformers
-pydantic==2.10.6
+transformers