{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yxig5CdZuHb9"
      },
      "source": [
        "# CountGD - Multimodal open-world object counting\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9wyM6J2HuHb-"
      },
      "source": [
        "## Setup\n",
        "\n",
        "The following cells will setup the runtime environment with the following\n",
        "\n",
        "- Mount Google Drive\n",
        "- Install dependencies for running the model\n",
        "- Load the model into memory"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jn061Tl8uHb-"
      },
      "source": [
        "### Mount Google Drive (if running on colab)\n",
        "\n",
        "The following bit of code will mount your Google Drive folder at `/content/drive`, allowing you to process files directly from it as well as store the results alongside it.\n",
        "\n",
        "Once you execute the next cell, you will be requested to share access with the notebook. Please follow the instructions on screen to do so.\n",
        "If you are not running this on colab, you will still be able to use the files available on your environment."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "collapsed": true,
        "id": "DkSUXqMPuHb-",
        "outputId": "6b82521e-3afd-4545-b13f-8cfea0975d95"
      },
      "outputs": [],
      "source": [
        "# Check if running colab\n",
        "import logging\n",
        "\n",
        "logging.basicConfig(\n",
        "  level=logging.INFO,\n",
        "  format='%(asctime)s %(levelname)-8s %(name)s %(message)s'\n",
        ")\n",
        "try:\n",
        "    import google.colab\n",
        "    RUNNING_IN_COLAB = True\n",
        "except:\n",
        "    RUNNING_IN_COLAB = False\n",
        "\n",
        "if RUNNING_IN_COLAB:\n",
        "    from google.colab import drive\n",
        "    drive.mount('/content/drive')\n",
        "\n",
        "from IPython.core.magic import register_cell_magic\n",
        "from IPython import get_ipython\n",
        "@register_cell_magic\n",
        "def skip_if(line, cell):\n",
        "    if eval(line):\n",
        "        return\n",
        "    get_ipython().run_cell(cell)\n",
        "\n",
        "\n",
        "%env RUNNING_IN_COLAB {RUNNING_IN_COLAB}\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kas5YtyluHb_"
      },
      "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)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "982Yiv5tuHb_",
        "outputId": "2f570d1a-c6cc-49c3-c336-1d784d33a169"
      },
      "outputs": [],
      "source": [
        "%%bash\n",
        "\n",
        "set -euxo pipefail\n",
        "\n",
        "if [ \"${RUNNING_IN_COLAB}\" == \"True\" ]; then\n",
        "  echo \"Downloading the repository...\"\n",
        "  if [ ! -d /content/countgd ]; then\n",
        "    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",
        "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",
        "  echo \"Running in $(pwd)\"\n",
        "fi\n",
        "\n",
        "# TODO check for gcc-11 or above\n",
        "\n",
        "# Install pip packages\n",
        "pip install --upgrade pip setuptools wheel\n",
        "pip install -r requirements.txt\n",
        "\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 test.py"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "58iD_HGnvcRJ",
        "outputId": "fe356a68-dced-4f6f-93cc-d83da2f84e28"
      },
      "outputs": [],
      "source": [
        "%cd {\"/content/countgd\" if RUNNING_IN_COLAB else '.'}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gH7A8zthuHb_"
      },
      "source": [
        "## Inference"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IspbBV0XuHb_"
      },
      "source": [
        "### Loading the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5nBT_HCUuHb_",
        "outputId": "95ceb6c6-bee8-4921-8bff-d28937045f78"
      },
      "outputs": [],
      "source": [
        "import app\n",
        "import importlib\n",
        "importlib.reload(app)\n",
        "from app import (\n",
        "    build_model_and_transforms,\n",
        "    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",
        "device = get_device()\n",
        "model, transform = build_model_and_transforms(args)\n",
        "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"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gfjraK3vuHb_"
      },
      "source": [
        "### Input / Output Utils\n",
        "\n",
        "Helper functions for reading / writing to zipfiles and csv"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 17,
      "metadata": {
        "id": "qg0g5B-fuHb_"
      },
      "outputs": [],
      "source": [
        "import io\n",
        "import csv\n",
        "from pathlib import Path\n",
        "from contextlib import contextmanager\n",
        "import zipfile\n",
        "import filetype\n",
        "from PIL import Image\n",
        "logger = logging.getLogger()\n",
        "\n",
        "def images_from_zipfile(p: Path):\n",
        "    if not zipfile.is_zipfile(p):\n",
        "        raise ValueError(f'{p} is not a zipfile!')\n",
        "\n",
        "    with zipfile.ZipFile(p, 'r') as zipf:\n",
        "        def process_entry(info: zipfile.ZipInfo):\n",
        "            with zipf.open(info) as f:\n",
        "                if not filetype.is_image(f):\n",
        "                    logger.debug(f'Skipping file - {info.filename} as it is not an image')\n",
        "                    return\n",
        "                # Try loading the file\n",
        "                try:\n",
        "                    with Image.open(f) as im:\n",
        "                        im.load()\n",
        "                        return (info.filename, im)\n",
        "                except:\n",
        "                    logger.exception(f'Error reading file {info.filename}')\n",
        "\n",
        "        num_files = sum(1 for info in zipf.infolist() if info.is_dir() == False)\n",
        "        logger.info(f'Found {num_files} file(s) in the zip')\n",
        "        yield from (process_entry(info) for info in zipf.infolist() if info.is_dir() == False)\n",
        "\n",
        "@contextmanager\n",
        "def zipfile_writer(p: Path):\n",
        "    with zipfile.ZipFile(p, 'w') as zipf:\n",
        "        def write_output(image, image_filename):\n",
        "            buf = io.BytesIO()\n",
        "            image.save(buf, 'PNG')\n",
        "            zipf.writestr(image_filename, buf.getvalue())\n",
        "        yield write_output\n",
        "\n",
        "@contextmanager\n",
        "def csvfile_writer(p: Path):\n",
        "    with p.open('w', newline='') as csvfile:\n",
        "        fieldnames = ['filename', 'count']\n",
        "        csv_writer = csv.DictWriter(csvfile, fieldnames = fieldnames)\n",
        "        csv_writer.writeheader()\n",
        "\n",
        "        yield csv_writer.writerow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rFXRk-_uuHb_"
      },
      "outputs": [],
      "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",
        "        return\n",
        "\n",
        "    if text == \"\":\n",
        "        logger.error('Please provide the object you would like to count')\n",
        "        return\n",
        "\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))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "TmqsSxrsuHb_"
      },
      "source": [
        "### Run\n",
        "\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."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "id": "ZaN918EkuHb_"
      },
      "outputs": [],
      "source": [
        "# @title ## Parameters { display-mode: \"form\", run: \"auto\" }\n",
        "# @markdown Set the following options to pass to the CountGD Model\n",
        "\n",
        "# @markdown ---\n",
        "# @markdown ### Enter a file path to a zip:\n",
        "zipfile_path = \"test_images.zip\" # @param {type:\"string\"}\n",
        "# @markdown\n",
        "# @markdown ### Which object would you like to count?\n",
        "prompt = \"strawberry\" # @param {type:\"string\"}\n",
        "# @markdown ---"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 66,
          "referenced_widgets": [
            "b14c910dd2594285bb4ad4740099e70c",
            "01631442369e43138c2c5c4a9fe38ceb",
            "ff84907ef88a431bab4bd3d1567cc42a"
          ]
        },
        "id": "fd-ShBCsuHb_",
        "outputId": "5b36bb90-ac6e-46fe-a853-ff11d43dd9f6"
      },
      "outputs": [],
      "source": [
        "import ipywidgets as widgets\n",
        "from IPython.display import display\n",
        "button = widgets.Button(description=\"Run\")\n",
        "\n",
        "def on_button_clicked(b):\n",
        "    # Display the message within the output widget.\n",
        "    process_zipfile(Path(zipfile_path), prompt)\n",
        "\n",
        "button.on_click(on_button_clicked)\n",
        "display(button)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "gfjraK3vuHb_"
      ],
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "env",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.7"
    },
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "01631442369e43138c2c5c4a9fe38ceb": {
          "model_module": "@jupyter-widgets/base",
          "model_module_version": "1.2.0",
          "model_name": "LayoutModel",
          "state": {
            "_model_module": "@jupyter-widgets/base",
            "_model_module_version": "1.2.0",
            "_model_name": "LayoutModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "LayoutView",
            "align_content": null,
            "align_items": null,
            "align_self": null,
            "border": null,
            "bottom": null,
            "display": null,
            "flex": null,
            "flex_flow": null,
            "grid_area": null,
            "grid_auto_columns": null,
            "grid_auto_flow": null,
            "grid_auto_rows": null,
            "grid_column": null,
            "grid_gap": null,
            "grid_row": null,
            "grid_template_areas": null,
            "grid_template_columns": null,
            "grid_template_rows": null,
            "height": null,
            "justify_content": null,
            "justify_items": null,
            "left": null,
            "margin": null,
            "max_height": null,
            "max_width": null,
            "min_height": null,
            "min_width": null,
            "object_fit": null,
            "object_position": null,
            "order": null,
            "overflow": null,
            "overflow_x": null,
            "overflow_y": null,
            "padding": null,
            "right": null,
            "top": null,
            "visibility": null,
            "width": null
          }
        },
        "b14c910dd2594285bb4ad4740099e70c": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ButtonModel",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ButtonView",
            "button_style": "",
            "description": "Run",
            "disabled": false,
            "icon": "",
            "layout": "IPY_MODEL_01631442369e43138c2c5c4a9fe38ceb",
            "style": "IPY_MODEL_ff84907ef88a431bab4bd3d1567cc42a",
            "tooltip": ""
          }
        },
        "ff84907ef88a431bab4bd3d1567cc42a": {
          "model_module": "@jupyter-widgets/controls",
          "model_module_version": "1.5.0",
          "model_name": "ButtonStyleModel",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ButtonStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "button_color": null,
            "font_weight": ""
          }
        }
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}