infestation_poc.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "could not find stamp in environment variables\n",
      "could not find stamp in environment variables\n",
      "could not find KUBEFLOW_API_BASIC_AUTH_TOKEN in environment variables\n",
      "could not find AI_RABBITMQ_CONNECTION_STRING in environment variables\n",
      "could not find AI_MONGO_DB_CONNECTION_STRING in environment variables\n",
      "could not find TARANIS_DB_CONNECTION_STRING in environment variables\n",
      "could not find GLOBAL_DB_CONNECTION_STRING in environment variables\n",
      "could not find AI_MONGO_DB_CONNECTION_STRING_NEW in environment variables\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import pandas as pd\n",
    "from skimage import io\n",
    "\n",
    "import numpy as np\n",
    "import PIL\n",
    "import torch\n",
    "\n",
    "from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor\n",
    "\n",
    "os.environ[\"NAMESPACE\"]=\"research\"\n",
    "os.environ[\"PROFILE\"]=\"local\"\n",
    "from agrobrain_util.runtime.evironment import RuntimeEnv\n",
    "\n",
    "DATA_DIR = \"data\"\n",
    "env = RuntimeEnv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_csv = pd.read_csv(os.path.join(DATA_DIR, \"1000_images_dataframe_for_tagging_methodology_task_2023_07_18.csv\"))\n",
    "\n",
    "example_image_id = int(images_csv['imageID'][0])\n",
    "im_path = env.download_image(int(example_image_id))\n",
    "\n",
    "example_wide_image_id = env.eti_api.get_matching_wide_images([5372412])[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "images_csv = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "ename": "RuntimeError",
     "evalue": "Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mRuntimeError\u001b[0m                              Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[5], line 7\u001b[0m\n\u001b[1;32m      4\u001b[0m device \u001b[39m=\u001b[39m \u001b[39m\"\u001b[39m\u001b[39mcuda\u001b[39m\u001b[39m\"\u001b[39m\n\u001b[1;32m      6\u001b[0m sam \u001b[39m=\u001b[39m sam_model_registry[model_type](checkpoint\u001b[39m=\u001b[39msam_checkpoint)\n\u001b[0;32m----> 7\u001b[0m sam\u001b[39m.\u001b[39;49mto(device)\n\u001b[1;32m      9\u001b[0m mask_generator \u001b[39m=\u001b[39m SamAutomaticMaskGenerator(\n\u001b[1;32m     10\u001b[0m     model\u001b[39m=\u001b[39msam,\n\u001b[1;32m     11\u001b[0m     points_per_side\u001b[39m=\u001b[39m\u001b[39m32\u001b[39m,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     16\u001b[0m     min_mask_region_area\u001b[39m=\u001b[39m\u001b[39m400\u001b[39m\n\u001b[1;32m     17\u001b[0m )\n\u001b[1;32m     19\u001b[0m image \u001b[39m=\u001b[39m np\u001b[39m.\u001b[39marray(PIL\u001b[39m.\u001b[39mImage\u001b[39m.\u001b[39mopen(im_path))\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1145\u001b[0m, in \u001b[0;36mModule.to\u001b[0;34m(self, *args, **kwargs)\u001b[0m\n\u001b[1;32m   1141\u001b[0m         \u001b[39mreturn\u001b[39;00m t\u001b[39m.\u001b[39mto(device, dtype \u001b[39mif\u001b[39;00m t\u001b[39m.\u001b[39mis_floating_point() \u001b[39mor\u001b[39;00m t\u001b[39m.\u001b[39mis_complex() \u001b[39melse\u001b[39;00m \u001b[39mNone\u001b[39;00m,\n\u001b[1;32m   1142\u001b[0m                     non_blocking, memory_format\u001b[39m=\u001b[39mconvert_to_format)\n\u001b[1;32m   1143\u001b[0m     \u001b[39mreturn\u001b[39;00m t\u001b[39m.\u001b[39mto(device, dtype \u001b[39mif\u001b[39;00m t\u001b[39m.\u001b[39mis_floating_point() \u001b[39mor\u001b[39;00m t\u001b[39m.\u001b[39mis_complex() \u001b[39melse\u001b[39;00m \u001b[39mNone\u001b[39;00m, non_blocking)\n\u001b[0;32m-> 1145\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_apply(convert)\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:797\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    795\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_apply\u001b[39m(\u001b[39mself\u001b[39m, fn):\n\u001b[1;32m    796\u001b[0m     \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mchildren():\n\u001b[0;32m--> 797\u001b[0m         module\u001b[39m.\u001b[39;49m_apply(fn)\n\u001b[1;32m    799\u001b[0m     \u001b[39mdef\u001b[39;00m \u001b[39mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    800\u001b[0m         \u001b[39mif\u001b[39;00m torch\u001b[39m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    801\u001b[0m             \u001b[39m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    802\u001b[0m             \u001b[39m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    807\u001b[0m             \u001b[39m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    808\u001b[0m             \u001b[39m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:797\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    795\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_apply\u001b[39m(\u001b[39mself\u001b[39m, fn):\n\u001b[1;32m    796\u001b[0m     \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mchildren():\n\u001b[0;32m--> 797\u001b[0m         module\u001b[39m.\u001b[39;49m_apply(fn)\n\u001b[1;32m    799\u001b[0m     \u001b[39mdef\u001b[39;00m \u001b[39mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    800\u001b[0m         \u001b[39mif\u001b[39;00m torch\u001b[39m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    801\u001b[0m             \u001b[39m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    802\u001b[0m             \u001b[39m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    807\u001b[0m             \u001b[39m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    808\u001b[0m             \u001b[39m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:797\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    795\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_apply\u001b[39m(\u001b[39mself\u001b[39m, fn):\n\u001b[1;32m    796\u001b[0m     \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mchildren():\n\u001b[0;32m--> 797\u001b[0m         module\u001b[39m.\u001b[39;49m_apply(fn)\n\u001b[1;32m    799\u001b[0m     \u001b[39mdef\u001b[39;00m \u001b[39mcompute_should_use_set_data\u001b[39m(tensor, tensor_applied):\n\u001b[1;32m    800\u001b[0m         \u001b[39mif\u001b[39;00m torch\u001b[39m.\u001b[39m_has_compatible_shallow_copy_type(tensor, tensor_applied):\n\u001b[1;32m    801\u001b[0m             \u001b[39m# If the new tensor has compatible tensor type as the existing tensor,\u001b[39;00m\n\u001b[1;32m    802\u001b[0m             \u001b[39m# the current behavior is to change the tensor in-place using `.data =`,\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    807\u001b[0m             \u001b[39m# global flag to let the user control whether they want the future\u001b[39;00m\n\u001b[1;32m    808\u001b[0m             \u001b[39m# behavior of overwriting the existing tensor or not.\u001b[39;00m\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:820\u001b[0m, in \u001b[0;36mModule._apply\u001b[0;34m(self, fn)\u001b[0m\n\u001b[1;32m    816\u001b[0m \u001b[39m# Tensors stored in modules are graph leaves, and we don't want to\u001b[39;00m\n\u001b[1;32m    817\u001b[0m \u001b[39m# track autograd history of `param_applied`, so we have to use\u001b[39;00m\n\u001b[1;32m    818\u001b[0m \u001b[39m# `with torch.no_grad():`\u001b[39;00m\n\u001b[1;32m    819\u001b[0m \u001b[39mwith\u001b[39;00m torch\u001b[39m.\u001b[39mno_grad():\n\u001b[0;32m--> 820\u001b[0m     param_applied \u001b[39m=\u001b[39m fn(param)\n\u001b[1;32m    821\u001b[0m should_use_set_data \u001b[39m=\u001b[39m compute_should_use_set_data(param, param_applied)\n\u001b[1;32m    822\u001b[0m \u001b[39mif\u001b[39;00m should_use_set_data:\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py:1143\u001b[0m, in \u001b[0;36mModule.to.<locals>.convert\u001b[0;34m(t)\u001b[0m\n\u001b[1;32m   1140\u001b[0m \u001b[39mif\u001b[39;00m convert_to_format \u001b[39mis\u001b[39;00m \u001b[39mnot\u001b[39;00m \u001b[39mNone\u001b[39;00m \u001b[39mand\u001b[39;00m t\u001b[39m.\u001b[39mdim() \u001b[39min\u001b[39;00m (\u001b[39m4\u001b[39m, \u001b[39m5\u001b[39m):\n\u001b[1;32m   1141\u001b[0m     \u001b[39mreturn\u001b[39;00m t\u001b[39m.\u001b[39mto(device, dtype \u001b[39mif\u001b[39;00m t\u001b[39m.\u001b[39mis_floating_point() \u001b[39mor\u001b[39;00m t\u001b[39m.\u001b[39mis_complex() \u001b[39melse\u001b[39;00m \u001b[39mNone\u001b[39;00m,\n\u001b[1;32m   1142\u001b[0m                 non_blocking, memory_format\u001b[39m=\u001b[39mconvert_to_format)\n\u001b[0;32m-> 1143\u001b[0m \u001b[39mreturn\u001b[39;00m t\u001b[39m.\u001b[39;49mto(device, dtype \u001b[39mif\u001b[39;49;00m t\u001b[39m.\u001b[39;49mis_floating_point() \u001b[39mor\u001b[39;49;00m t\u001b[39m.\u001b[39;49mis_complex() \u001b[39melse\u001b[39;49;00m \u001b[39mNone\u001b[39;49;00m, non_blocking)\n",
      "File \u001b[0;32m/usr/local/lib/python3.8/dist-packages/torch/cuda/__init__.py:247\u001b[0m, in \u001b[0;36m_lazy_init\u001b[0;34m()\u001b[0m\n\u001b[1;32m    245\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39m'\u001b[39m\u001b[39mCUDA_MODULE_LOADING\u001b[39m\u001b[39m'\u001b[39m \u001b[39mnot\u001b[39;00m \u001b[39min\u001b[39;00m os\u001b[39m.\u001b[39menviron:\n\u001b[1;32m    246\u001b[0m     os\u001b[39m.\u001b[39menviron[\u001b[39m'\u001b[39m\u001b[39mCUDA_MODULE_LOADING\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m=\u001b[39m \u001b[39m'\u001b[39m\u001b[39mLAZY\u001b[39m\u001b[39m'\u001b[39m\n\u001b[0;32m--> 247\u001b[0m torch\u001b[39m.\u001b[39;49m_C\u001b[39m.\u001b[39;49m_cuda_init()\n\u001b[1;32m    248\u001b[0m \u001b[39m# Some of the queued calls may reentrantly call _lazy_init();\u001b[39;00m\n\u001b[1;32m    249\u001b[0m \u001b[39m# we need to just return without initializing in that case.\u001b[39;00m\n\u001b[1;32m    250\u001b[0m \u001b[39m# However, we must not let any *other* threads in!\u001b[39;00m\n\u001b[1;32m    251\u001b[0m _tls\u001b[39m.\u001b[39mis_initializing \u001b[39m=\u001b[39m \u001b[39mTrue\u001b[39;00m\n",
      "\u001b[0;31mRuntimeError\u001b[0m: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx"
     ]
    }
   ],
   "source": [
    "sam_checkpoint = os.path.join(\"models\", \"sam_vit_b_01ec64.pth\")\n",
    "# \"path/to/ckpt.pth\"\n",
    "model_type = \"vit_b\"\n",
    "device = \"cuda\"\n",
    "\n",
    "sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)\n",
    "sam.to(device)\n",
    "\n",
    "mask_generator = SamAutomaticMaskGenerator(\n",
    "    model=sam,\n",
    "    points_per_side=32,\n",
    "    pred_iou_thresh=0.9,\n",
    "    stability_score_thresh=0.92,\n",
    "    crop_n_layers=1,\n",
    "    crop_n_points_downscale_factor=2,\n",
    "    min_mask_region_area=400\n",
    ")\n",
    "\n",
    "image = np.array(PIL.Image.open(im_path))\n",
    "masks = mask_generator.generate(image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SHOW IMAGE AND MASK\n",
    "\n",
    "fig, axes = plt.subplots(nrows=1, ncols=2)\n",
    "axes[0].imshow(image, extent=[0, image.shape[1], 0, image.shape[0]])\n",
    "axes[1].imshow(masks, extent=[0, masks.shape[1], 0, masks.shape[0]])\n",
    "\n",
    "\n",
    "for ax in axes:\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "fig.set_size_inches(10, 5)\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.10"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}