Export box coordinates in demo notebook

app.py

@@ -176,11 +176,27 @@ 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))
-    det_map[(h * boxes[:, 1]).astype(int), (w * boxes[:, 0]).astype(int)] = 1
+    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 = ndimage.gaussian_filter(
         det_map, sigma=(w // 200, w // 200), order=0
     )

notebooks/demo.ipynb

@@ -6,7 +6,7 @@
         "id": "yxig5CdZuHb9"
       },
       "source": [
-        "# CountGD - Multimodela open-world object counting\n",
+        "# CountGD - Multimodal open-world object counting\n",
         "\n"
       ]
     },
@@ -89,7 +89,13 @@
       "source": [
         "### Install Dependencies\n",
         "\n",
-        "The environment will be setup with the code, models and required dependencies."
+        "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)"
       ]
     },
     {
@@ -114,8 +120,10 @@
         "    git clone \"https://huggingface.co/spaces/nikigoli/countgd\" /content/countgd\n",
         "  fi\n",
         "  cd /content/countgd\n",
-        "  git fetch origin refs/pr/5:refs/remotes/origin/pr/5\n",
-        "  git checkout pr/5\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",
         "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",
@@ -128,11 +136,20 @@
         "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",
-        "python3 setup.py build\n",
-        "pip install .\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 test.py"
       ]
     },
@@ -189,6 +206,7 @@
         "    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",
@@ -197,7 +215,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), generate_heatmap(image, boxes))\n"
+        "get_output = lambda image, boxes: (len(boxes), get_xy_from_boxes(image, boxes), generate_heatmap(image, boxes))\n"
       ]
     },
     {
@@ -271,7 +289,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 15,
+      "execution_count": null,
       "metadata": {
         "id": "rFXRk-_uuHb_"
       },
@@ -279,6 +297,17 @@
       "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",
@@ -290,14 +319,24 @@
         "\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",
-        "            boxes, _ = run(im, text)\n",
-        "            count, heatmap  = get_output(im, boxes)\n",
-        "            write_row({'filename': filename, 'count': count})\n",
-        "            add_to_zip(heatmap, filename)"
+        "            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))"
       ]
     },
     {
@@ -310,6 +349,8 @@
         "\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."