SAMTrack_inference.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "a9d0125c361a426e9d5f8ac311f061c0": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_a9db6b75095b4c979f23f96634cf7992",
              "IPY_MODEL_24df4c73fa064e0f969ae40552335fe7",
              "IPY_MODEL_969bf22f1b9e46bc98ae9fe7fce65086"
            ],
            "layout": "IPY_MODEL_9dfc41ec4a984bbb9820dd6f07003535"
          }
        },
        "a9db6b75095b4c979f23f96634cf7992": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_3aeeda7db5f8493b903ff06453156670",
            "placeholder": "​",
            "style": "IPY_MODEL_0a4e220d8cdc44b4b62fdac6ca14e30c",
            "value": "100%"
          }
        },
        "24df4c73fa064e0f969ae40552335fe7": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_b109b007ea2d4fe29f3b4038961080a9",
            "max": 51,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_22a376ea2da845db978aa6452e5a0c16",
            "value": 51
          }
        },
        "969bf22f1b9e46bc98ae9fe7fce65086": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_1789f4e986da4dcc882345a7360fccb4",
            "placeholder": "​",
            "style": "IPY_MODEL_7ba4c479c09b40e8bfbe9d68de309d05",
            "value": " 51/51 [00:29&lt;00:00,  1.19it/s]"
          }
        },
        "9dfc41ec4a984bbb9820dd6f07003535": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "3aeeda7db5f8493b903ff06453156670": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "0a4e220d8cdc44b4b62fdac6ca14e30c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b109b007ea2d4fe29f3b4038961080a9": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "22a376ea2da845db978aa6452e5a0c16": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "1789f4e986da4dcc882345a7360fccb4": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "7ba4c479c09b40e8bfbe9d68de309d05": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "b8246408a9bc4d519b738cd6af2db010": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HBoxModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HBoxView",
            "box_style": "",
            "children": [
              "IPY_MODEL_04e9b3573c374f76be694e4f4a326b95",
              "IPY_MODEL_039a6f2543f044028cfd6d3dafcb1696",
              "IPY_MODEL_5bbb7c1cd66a45a59870187f477bb6cb"
            ],
            "layout": "IPY_MODEL_3a51c99a82174dab990205d6c0eec4b2"
          }
        },
        "04e9b3573c374f76be694e4f4a326b95": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_f1e66c47264f4c31ba5a1afe74f3e7b2",
            "placeholder": "​",
            "style": "IPY_MODEL_bf418989bcce4b16a8aaa8d9568a3cb1",
            "value": "100%"
          }
        },
        "039a6f2543f044028cfd6d3dafcb1696": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "FloatProgressModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "ProgressView",
            "bar_style": "success",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_308b96514d91468fa86746da02629b74",
            "max": 51,
            "min": 0,
            "orientation": "horizontal",
            "style": "IPY_MODEL_be5a771d7a774e8aa3448c3e474bec76",
            "value": 51
          }
        },
        "5bbb7c1cd66a45a59870187f477bb6cb": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "model_module_version": "1.5.0",
          "state": {
            "_dom_classes": [],
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "HTMLModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/controls",
            "_view_module_version": "1.5.0",
            "_view_name": "HTMLView",
            "description": "",
            "description_tooltip": null,
            "layout": "IPY_MODEL_47794863dd2a4280a9cdd460e2415df5",
            "placeholder": "​",
            "style": "IPY_MODEL_1f1f7564317841cb8fe5e80b36cd61bd",
            "value": " 51/51 [00:00&lt;00:00, 94.79it/s]"
          }
        },
        "3a51c99a82174dab990205d6c0eec4b2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "f1e66c47264f4c31ba5a1afe74f3e7b2": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "bf418989bcce4b16a8aaa8d9568a3cb1": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        },
        "308b96514d91468fa86746da02629b74": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "be5a771d7a774e8aa3448c3e474bec76": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "ProgressStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "bar_color": null,
            "description_width": ""
          }
        },
        "47794863dd2a4280a9cdd460e2415df5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "model_module_version": "1.2.0",
          "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
          }
        },
        "1f1f7564317841cb8fe5e80b36cd61bd": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "model_module_version": "1.5.0",
          "state": {
            "_model_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_model_name": "DescriptionStyleModel",
            "_view_count": null,
            "_view_module": "@jupyter-widgets/base",
            "_view_module_version": "1.2.0",
            "_view_name": "StyleView",
            "description_width": ""
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "SAMTrack inference notebook by [Alex Spirin](https://github.com/Sxela)\n",
        "\n",
        "The main difference from the [official repo](https://github.com/z-x-yang/Segment-and-Track-Anything) is that it exports masks separately for easier compositing.\n",
        "\n",
        "License: [AGPL](https://github.com/Sxela/Segment-and-Track-Anything-CLI/blob/main/LICENSE.txt)\n",
        "## Local Installation prerequsites\n",
        "Pre-built groundingDino and spatial-correlation-sampler binaries for local Win11/cuda11.8/python3.10 are downloaded automatically.\n",
        "\n",
        "If the binaries didn't work for you:\n",
        "\n",
        "- Get [MSVC Build tools](https://aka.ms/vs/17/release/vs_BuildTools.exe) and install the local c++ dev kit\n",
        "-  Get [latest nVidia CUDA toolkit](https://developer.nvidia.com/cuda-downloads?target_os=Windows&target_arch=x86_64&target_version=11&target_type=exe_local) or at least [11.8+](https://developer.nvidia.com/cuda-11-8-0-download-archive?target_os=Windows&target_arch=x86_64&target_version=11&target_type=exe_local) and install it. Don't forget to remove older versions."
      ],
      "metadata": {
        "id": "LTud6YoSMOdB"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Install SAMTrack-CLI\n",
        "#@markdown originally from https://github.com/z-x-yang/Segment-and-Track-Anything \\\n",
        "#@markdown Restart the notebook after install.\n",
        "#https://stackoverflow.com/questions/64261546/how-to-solve-error-microsoft-visual-c-14-0-or-greater-is-required-when-inst\n",
        "import os, platform\n",
        "try:\n",
        "  #cd to root if root dir defined\n",
        "  os.chdir(root_dir)\n",
        "except:\n",
        "  root_dir = os.getcwd()\n",
        "\n",
        "!git clone https://github.com/Sxela/Segment-and-Track-Anything-CLI\n",
        "os.chdir(os.path.join(root_dir,'Segment-and-Track-Anything-CLI'))\n",
        "\n",
        "!python -m pip install -e ./sam\n",
        "if platform.system() == 'Linux':\n",
        "  !python -m pip install -e git+https://github.com/IDEA-Research/GroundingDINO.git@main#egg=GroundingDINO\n",
        "else:\n",
        "  os.makedirs('./src', exist_ok=True)\n",
        "  !git clone https://github.com/IDEA-Research/GroundingDINO \"{os.path.join(root_dir,'Segment-and-Track-Anything-CLI')}/src/GroundingDINO\"\n",
        "  !python -m pip install -r \"{os.path.join(root_dir,'Segment-and-Track-Anything-CLI')}/src/GroundingDINO/requirements.txt\"\n",
        "!python -m pip install numpy opencv-python pycocotools matplotlib Pillow scikit-image\n",
        "!python -m pip install gdown\n",
        "!python -m pip install wget\n",
        "!git clone https://github.com/ClementPinard/Pytorch-Correlation-extension.git\n",
        "if platform.system() == 'Linux':\n",
        "  !python -m pip install -e ./Pytorch-Correlation-extension\n",
        "else:\n",
        "  !python -m pip install -r ./Pytorch-Correlation-extension/requirements.txt\n",
        "\n",
        "os.chdir(os.path.join(root_dir,'Segment-and-Track-Anything-CLI'))\n",
        "os.makedirs(os.path.join(root_dir,'Segment-and-Track-Anything-CLI', 'ckpt'), exist_ok=True)\n",
        "\n",
        "import gdown\n",
        "# download aot-ckpt\n",
        "if not os.path.exists('./ckpt/R50_DeAOTL_PRE_YTB_DAV.pth'):\n",
        "  gdown.download(id='1QoChMkTVxdYZ_eBlZhK2acq9KMQZccPJ', output='./ckpt/R50_DeAOTL_PRE_YTB_DAV.pth')\n",
        "\n",
        "import wget\n",
        "# download sam-ckpt\n",
        "if not os.path.exists('./ckpt/sam_vit_b_01ec64.pth'):\n",
        "  wget.download(\"https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth\",\n",
        "              \"ckpt/\")\n",
        "\n",
        "if not os.path.exists('./ckpt/groundingdino_swint_ogc.pth'):\n",
        "# download grounding-dino ckpt\n",
        "  wget.download(\"https://huggingface.co/ShilongLiu/GroundingDINO/resolve/main/groundingdino_swint_ogc.pth\",\n",
        "              \"ckpt/\")\n",
        "\n",
        "import wget\n",
        "import zipfile\n",
        "\n",
        "if platform.system() != 'Linux':\n",
        "  #download prebuilt binaries for cuda 11.8, torch 2, python 3.10, win11\n",
        "  if not os.path.exists('./site-packages.zip'):\n",
        "    wget.download(\"https://raw.githubusercontent.com/Sxela/Segment-and-Track-Anything-CLI/main/site-packages.zip\",\n",
        "                \"./site-packages.zip\")\n",
        "\n",
        "  with zipfile.ZipFile(\"site-packages.zip\", 'r') as zip_ref:\n",
        "          zip_ref.extractall(f'{root_dir}/env/Lib/')"
      ],
      "metadata": {
        "cellView": "form",
        "id": "5uXN8KDDMpV_"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Detection setup\n",
        "#@markdown Use this cell to tweak detection settings, that will be later used on the whole video.\n",
        "#@markdown Run this cell to get detection preview.\\\n",
        "#@markdown Code mostly taken from https://github.com/z-x-yang/Segment-and-Track-Anything/blob/main/demo_instseg.ipynb\n",
        "import os, pathlib, shutil, sys, subprocess\n",
        "from glob import glob\n",
        "try:\n",
        "  #cd to root if root dir defined\n",
        "  os.chdir(root_dir)\n",
        "except:\n",
        "  root_dir = os.getcwd()\n",
        "\n",
        "os.chdir(os.path.join(root_dir,'Segment-and-Track-Anything-CLI'))\n",
        "\n",
        "#(c) Alex Spirin 2023\n",
        "\n",
        "import hashlib\n",
        "# We use input file hashes to automate video extraction\n",
        "#\n",
        "def generate_file_hash(input_file):\n",
        "    # Get file name and metadata\n",
        "    file_name = os.path.basename(input_file)\n",
        "    file_size = os.path.getsize(input_file)\n",
        "    creation_time = os.path.getctime(input_file)\n",
        "\n",
        "    # Generate hash\n",
        "    hasher = hashlib.sha256()\n",
        "    hasher.update(file_name.encode('utf-8'))\n",
        "    hasher.update(str(file_size).encode('utf-8'))\n",
        "    hasher.update(str(creation_time).encode('utf-8'))\n",
        "    file_hash = hasher.hexdigest()\n",
        "\n",
        "    return file_hash\n",
        "\n",
        "def createPath(filepath):\n",
        "    os.makedirs(filepath, exist_ok=True)\n",
        "\n",
        "\n",
        "def extractFrames(video_path, output_path, nth_frame, start_frame, end_frame):\n",
        "  createPath(output_path)\n",
        "  print(f\"Exporting Video Frames (1 every {nth_frame})...\")\n",
        "  try:\n",
        "    for f in [o.replace('\\\\','/') for o in glob(output_path+'/*.jpg')]:\n",
        "    # for f in pathlib.Path(f'{output_path}').glob('*.jpg'):\n",
        "      pathlib.Path(f).unlink()\n",
        "  except:\n",
        "    print('error deleting frame ', f)\n",
        "  # vf = f'select=not(mod(n\\\\,{nth_frame}))'\n",
        "  vf = f'select=between(n\\\\,{start_frame}\\\\,{end_frame}) , select=not(mod(n\\\\,{nth_frame}))'\n",
        "  if os.path.exists(video_path):\n",
        "    try:\n",
        "        # subprocess.run(['ffmpeg', '-i', f'{video_path}', '-vf', f'{vf}', '-vsync', 'vfr', '-q:v', '2', '-loglevel', 'error', '-stats', f'{output_path}/%06d.jpg'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "\n",
        "        subprocess.run(['ffmpeg', '-i', f'{video_path}', '-vf', f'{vf}', '-vsync', 'vfr', '-q:v', '2', '-loglevel', 'error', '-stats', f'{output_path}/%06d.jpg'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "    except:\n",
        "        subprocess.run([f'{root_dir}/ffmpeg.exe', '-i', f'{video_path}', '-vf', f'{vf}', '-vsync', 'vfr', '-q:v', '2', '-loglevel', 'error', '-stats', f'{output_path}/%06d.jpg'], stdout=subprocess.PIPE).stdout.decode('utf-8')\n",
        "\n",
        "  else:\n",
        "    sys.exit(f'\\nERROR!\\n\\nVideo not found: {video_path}.\\nPlease check your video path.\\n')\n",
        "\n",
        "\n",
        "class FrameDataset():\n",
        "  def __init__(self, source_path, outdir_prefix, videoframes_root):\n",
        "    self.frame_paths = None\n",
        "    image_extenstions = ['jpeg', 'jpg', 'png', 'tiff', 'bmp', 'webp']\n",
        "\n",
        "    if not os.path.exists(source_path):\n",
        "      if len(glob(source_path))>0:\n",
        "        self.frame_paths = sorted(glob(source_path))\n",
        "      else:\n",
        "        raise Exception(f'Frame source for {outdir_prefix} not found at {source_path}\\nPlease specify an existing source path.')\n",
        "    if os.path.exists(source_path):\n",
        "      if os.path.isfile(source_path):\n",
        "        if os.path.splitext(source_path)[1][1:].lower() in image_extenstions:\n",
        "          self.frame_paths = [source_path]\n",
        "        hash = generate_file_hash(source_path)[:10]\n",
        "        out_path = os.path.join(videoframes_root, outdir_prefix+'_'+hash)\n",
        "\n",
        "        extractFrames(source_path, out_path,\n",
        "                        nth_frame=1, start_frame=0, end_frame=999999999)\n",
        "        self.frame_paths = glob(os.path.join(out_path, '*.*'))\n",
        "        if len(self.frame_paths)<1:\n",
        "            raise Exception(f'Couldn`t extract frames from {source_path}\\nPlease specify an existing source path.')\n",
        "      elif os.path.isdir(source_path):\n",
        "        self.frame_paths = glob(os.path.join(source_path, '*.*'))\n",
        "        if len(self.frame_paths)<1:\n",
        "          raise Exception(f'Found 0 frames in {source_path}\\nPlease specify an existing source path.')\n",
        "    extensions = []\n",
        "    if self.frame_paths is not None:\n",
        "      for f in self.frame_paths:\n",
        "            ext = os.path.splitext(f)[1][1:]\n",
        "            if ext not in image_extenstions:\n",
        "              raise Exception(f'Found non-image file extension: {ext} in {source_path}. Please provide a folder with image files of the same extension, or specify a glob pattern.')\n",
        "            if not ext in extensions:\n",
        "              extensions+=[ext]\n",
        "            if len(extensions)>1:\n",
        "              raise Exception(f'Found multiple file extensions: {extensions} in {source_path}. Please provide a folder with image files of the same extension, or specify a glob pattern.')\n",
        "\n",
        "      self.frame_paths = sorted(self.frame_paths)\n",
        "\n",
        "    else: raise Exception(f'Frame source for {outdir_prefix} not found at {source_path}\\nPlease specify an existing source path.')\n",
        "    print(f'Found {len(self.frame_paths)} frames at {source_path}')\n",
        "\n",
        "  def __getitem__(self, idx):\n",
        "    idx = min(idx, len(self.frame_paths)-1)\n",
        "    return self.frame_paths[idx]\n",
        "\n",
        "  def __len__(self):\n",
        "    return len(self.frame_paths)\n",
        "\n",
        "# mostly taken from https://github.com/z-x-yang/Segment-and-Track-Anything/blob/main/demo_instseg.ipynb\n",
        "\n",
        "import os\n",
        "import cv2\n",
        "from SegTracker import SegTracker\n",
        "from model_args import aot_args,sam_args,segtracker_args\n",
        "from PIL import Image\n",
        "from aot_tracker import _palette\n",
        "import numpy as np\n",
        "import torch\n",
        "import imageio\n",
        "import matplotlib.pyplot as plt\n",
        "from scipy.ndimage import binary_dilation\n",
        "\n",
        "import gc\n",
        "def save_prediction(pred_mask,output_dir,file_name):\n",
        "    save_mask = Image.fromarray(pred_mask.astype(np.uint8))\n",
        "    save_mask = save_mask.convert(mode='P')\n",
        "    save_mask.putpalette(_palette)\n",
        "    save_mask.save(os.path.join(output_dir,file_name))\n",
        "def colorize_mask(pred_mask):\n",
        "    save_mask = Image.fromarray(pred_mask.astype(np.uint8))\n",
        "    save_mask = save_mask.convert(mode='P')\n",
        "    save_mask.putpalette(_palette)\n",
        "    save_mask = save_mask.convert(mode='RGB')\n",
        "    return np.array(save_mask)\n",
        "def draw_mask(img, mask, alpha=0.7, id_countour=False):\n",
        "    img_mask = np.zeros_like(img)\n",
        "    img_mask = img\n",
        "    if id_countour:\n",
        "        # very slow ~ 1s per image\n",
        "        obj_ids = np.unique(mask)\n",
        "        obj_ids = obj_ids[obj_ids!=0]\n",
        "\n",
        "        for id in obj_ids:\n",
        "            # Overlay color on  binary mask\n",
        "            if id <= 255:\n",
        "                color = _palette[id*3:id*3+3]\n",
        "            else:\n",
        "                color = [0,0,0]\n",
        "            foreground = img * (1-alpha) + np.ones_like(img) * alpha * np.array(color)\n",
        "            binary_mask = (mask == id)\n",
        "\n",
        "            # Compose image\n",
        "            img_mask[binary_mask] = foreground[binary_mask]\n",
        "\n",
        "            countours = binary_dilation(binary_mask,iterations=1) ^ binary_mask\n",
        "            img_mask[countours, :] = 0\n",
        "    else:\n",
        "        binary_mask = (mask!=0)\n",
        "        countours = binary_dilation(binary_mask,iterations=1) ^ binary_mask\n",
        "        foreground = img*(1-alpha)+colorize_mask(mask)*alpha\n",
        "        img_mask[binary_mask] = foreground[binary_mask]\n",
        "        img_mask[countours,:] = 0\n",
        "\n",
        "    return img_mask.astype(img.dtype)\n",
        "\n",
        "video_path = '/content/chess girl.mov' #@param {'type':'string'}\n",
        "video_name = video_path.replace('\\\\','/').split('/')[-1]\n",
        "io_args = {\n",
        "    'input_video': video_path,\n",
        "    'output_mask_dir': f'./assets/{video_name}_masks', # save pred masks\n",
        "    'output_video': f'./assets/{video_name}_seg.mp4', # mask+frame vizualization, mp4 or avi, else the same as input video\n",
        "    'output_gif': f'./assets/{video_name}_seg.gif', # mask visualization\n",
        "}\n",
        "prefix = ''\n",
        "try:\n",
        "  videoframes_root = f'{batchFolder}/videoFrames'\n",
        "except:\n",
        "  videoframes_root = f'{root_dir}/videoFrames'\n",
        "\n",
        "frames = FrameDataset(video_path, outdir_prefix=prefix, videoframes_root=videoframes_root)\n",
        "\n",
        "# choose good parameters in sam_args based on the first frame segmentation result\n",
        "# other arguments can be modified in model_args.py\n",
        "# note the object number limit is 255 by default, which requires < 10GB GPU memory with amp\n",
        "sam_args['generator_args'] = {\n",
        "        'points_per_side': 60,\n",
        "        'pred_iou_thresh': 0.8,\n",
        "        'stability_score_thresh': 0.9,\n",
        "        'crop_n_layers': 1,\n",
        "        'crop_n_points_downscale_factor': 2,\n",
        "        'min_mask_region_area': 200,\n",
        "    }\n",
        "\n",
        "# Set Text args\n",
        "'''\n",
        "parameter:\n",
        "    grounding_caption: Text prompt to detect objects in key-frames\n",
        "    box_threshold: threshold for box\n",
        "    text_threshold: threshold for label(text)\n",
        "    box_size_threshold: If the size ratio between the box and the frame is larger than the box_size_threshold, the box will be ignored. This is used to filter out large boxes.\n",
        "    reset_image: reset the image embeddings for SAM\n",
        "'''\n",
        "frame_number = 0  #@param {'type':'number'}\n",
        "frame_number = int(frame_number)\n",
        "#@markdown Text prompt to detect objects in key-frames\n",
        "grounding_caption = \"person\" #@param {'type':'string'}\n",
        "#@markdown Box detection confidence threshold\n",
        "box_threshold = 0.3 #@param {'type':'number'}\n",
        "#@markdown Text confidence threshold\n",
        "text_threshold = 0.3 #@param {'type':'number'}\n",
        "#@markdown Box to Image ratio threshold (with box_size_threshold = 0.8 detections over 80% of the image will be ignored)\n",
        "box_size_threshold = 1 #@param {'type':'number'}\n",
        "\n",
        "reset_image = True\n",
        "\n",
        "frame_idx = 0\n",
        "segtracker = SegTracker(segtracker_args,sam_args,aot_args)\n",
        "segtracker.restart_tracker()\n",
        "\n",
        "with torch.cuda.amp.autocast():\n",
        "    frame = cv2.imread(frames[frame_number])\n",
        "    frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
        "    pred_mask, annotated_frame = segtracker.detect_and_seg(frame, grounding_caption, box_threshold, text_threshold,\n",
        "                                                           box_size_threshold, reset_image=reset_image)\n",
        "    torch.cuda.empty_cache()\n",
        "    obj_ids = np.unique(pred_mask)\n",
        "    obj_ids = obj_ids[obj_ids!=0]\n",
        "    print(\"processed frame {}, obj_num {}\".format(frame_idx,len(obj_ids)),end='\\n')\n",
        "    init_res = draw_mask(annotated_frame, pred_mask,id_countour=False)\n",
        "    plt.figure(figsize=(10,10))\n",
        "    plt.axis('off')\n",
        "    plt.imshow(init_res)\n",
        "    plt.show()\n",
        "    plt.figure(figsize=(10,10))\n",
        "    plt.axis('off')\n",
        "    plt.imshow(colorize_mask(pred_mask))\n",
        "    plt.show()\n",
        "\n",
        "    del segtracker\n",
        "    torch.cuda.empty_cache()\n",
        "    gc.collect()"
      ],
      "metadata": {
        "cellView": "form",
        "id": "UxbvG8XeNqB3"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#@title Mask whole video.\n",
        "use_cli = False #@param {'type':'boolean'}\n",
        "import subprocess\n",
        "start_frame = 0 #@param {'type':'number'}\n",
        "end_frame = 50 #@param {'type':'number'}\n",
        "#@markdown The interval to run SAM to segment new objects\n",
        "sam_gap = 50 #@param {'type':'number'}\n",
        "#@markdown minimal mask area to add a new mask as a new object\n",
        "min_area = 200  #@param {'type':'number'}\n",
        "#@markdown maximal object number to track in a video\n",
        "max_obj_num = 255 #@param {'type':'number'}\n",
        "#@markdown the area of a new object in the background should > 80%\n",
        "min_new_obj_iou = 0.8 #@param {'type':'number'}\n",
        "save_separate_masks = True\n",
        "save_joint_mask = False #@param {'type':'boolean'}\n",
        "save_mask = save_joint_mask\n",
        "save_video = False #@param {'type':'boolean'}\n",
        "save_gif = False #@param {'type':'boolean'}\n",
        "# grounding_caption\n",
        "# box_threshold\n",
        "# text_threshold\n",
        "# box_size_threshold\n",
        "# video_path\n",
        "output_multimask_dir = os.path.join(videoframes_root, f'{generate_file_hash(video_path)[:10]}_masks')\n",
        "if use_cli:\n",
        "  def run_command(cmd, cwd='./'):\n",
        "      with subprocess.Popen(cmd, stdout=subprocess.PIPE, bufsize=1, universal_newlines=True) as p:\n",
        "          while True:\n",
        "              line = p.stdout.readline()\n",
        "              if not line:\n",
        "                  break\n",
        "              print(line)\n",
        "          exit_code = p.poll()\n",
        "      return exit_code\n",
        "\n",
        "  # !python /content/Segment-and-Track-Anything/run.py\\\n",
        "  #  --video_path /content/SaveInsta.App_-_3067564057762969265_1317509610.mp4\\\n",
        "  #  --save_separate_masks --outdir /content/out/\n",
        "\n",
        "\n",
        "  cmd = ['python', 'run.py','--video_path', video_path, '--save_separate_masks', '--outdir', output_multimask_dir,\n",
        "        '--caption', grounding_caption, '--box_threshold', box_threshold, '--text_threshold', text_threshold, '--box_size_threshold', box_size_threshold,\n",
        "        '--sam_gap', sam_gap, '--min_area', min_area, '--max_obj_num', max_obj_num, '--min_new_obj_iou',min_new_obj_iou]\n",
        "  cmd = [str(o) for o in cmd]\n",
        "  returncode = run_command(cmd, cwd=os.path.join(root_dir,'Segment-and-Track-Anything-CLI'))\n",
        "  if process.returncode != 0:\n",
        "    raise RuntimeError(returncode)\n",
        "  else:\n",
        "    print(f\"The video is ready and saved to {output_multimask_dir}\")\n",
        "else:\n",
        "  os.makedirs('./debug/seg_result', exist_ok=True)\n",
        "  os.makedirs('./debug/aot_result', exist_ok=True)\n",
        "  segtracker_args = {\n",
        "    'sam_gap': sam_gap,\n",
        "    'min_area': min_area,\n",
        "    'max_obj_num': max_obj_num,\n",
        "    'min_new_obj_iou': min_new_obj_iou\n",
        "  }\n",
        "\n",
        "  if save_mask:\n",
        "    output_dir = io_args['output_mask_dir']\n",
        "    os.makedirs(output_dir, exist_ok=True)\n",
        "  pred_list = []\n",
        "  masked_pred_list = []\n",
        "\n",
        "  segtracker = SegTracker(segtracker_args, sam_args, aot_args)\n",
        "  segtracker.restart_tracker()\n",
        "  from tqdm.notebook import tqdm, trange\n",
        "  if start_frame == 0 and end_frame == 0:\n",
        "    frame_range = trange(len(frames))\n",
        "  else:\n",
        "    frame_range = trange(start_frame, end_frame+1)\n",
        "  for frame_idx in frame_range:\n",
        "    frame = cv2.imread(frames[frame_idx])\n",
        "    frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
        "    if frame_idx == start_frame:\n",
        "      pred_mask, _ = segtracker.detect_and_seg(frame, grounding_caption, box_threshold, text_threshold, box_size_threshold, reset_image)\n",
        "      torch.cuda.empty_cache()\n",
        "      gc.collect()\n",
        "      segtracker.add_reference(frame, pred_mask)\n",
        "    elif ((frame_idx-start_frame) % sam_gap) == 0:\n",
        "      seg_mask, _ = segtracker.detect_and_seg(frame, grounding_caption, box_threshold, text_threshold,\n",
        "                                                    box_size_threshold, reset_image)\n",
        "      # save_prediction(seg_mask, './debug/seg_result', str(frame_idx)+'.png')\n",
        "      torch.cuda.empty_cache()\n",
        "      gc.collect()\n",
        "      track_mask = segtracker.track(frame)\n",
        "      # save_prediction(track_mask, './debug/aot_result', str(frame_idx)+'.png')\n",
        "\n",
        "      # find new objects, and update tracker with new objects\n",
        "      new_obj_mask = segtracker.find_new_objs(track_mask, seg_mask)\n",
        "      if np.sum(new_obj_mask > 0) >  frame.shape[0] * frame.shape[1] * 0.4:\n",
        "        new_obj_mask = np.zeros_like(new_obj_mask)\n",
        "      if save_mask: save_prediction(new_obj_mask,output_dir,str(frame_idx)+'_new.png')\n",
        "      pred_mask = track_mask + new_obj_mask\n",
        "      segtracker.add_reference(frame, pred_mask)\n",
        "    else:\n",
        "      pred_mask = segtracker.track(frame,update_memory=True)\n",
        "    torch.cuda.empty_cache()\n",
        "    gc.collect()\n",
        "\n",
        "    if save_mask: save_prediction(pred_mask,output_dir,str(frame_idx)+'.png')\n",
        "\n",
        "    pred_list.append(pred_mask)\n",
        "\n",
        "    print(\"processed frame {}, obj_num {}\".format(frame_idx,segtracker.get_obj_num()),end='\\r')\n",
        "\n",
        "\n",
        "  if  save_video:\n",
        "  # draw pred mask on frame and save as a video\n",
        "    cap = cv2.VideoCapture(io_args['input_video'])\n",
        "    fps = cap.get(cv2.CAP_PROP_FPS)\n",
        "    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))\n",
        "    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))\n",
        "    num_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n",
        "\n",
        "    if io_args['input_video'][-3:]=='mp4':\n",
        "        fourcc =  cv2.VideoWriter_fourcc(*\"mp4v\")\n",
        "    elif io_args['input_video'][-3:] == 'avi':\n",
        "        fourcc =  cv2.VideoWriter_fourcc(*\"MJPG\")\n",
        "        # fourcc = cv2.VideoWriter_fourcc(*\"XVID\")\n",
        "    else:\n",
        "        fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))\n",
        "    out = cv2.VideoWriter(io_args['output_video'], fourcc, fps, (width, height))\n",
        "\n",
        "    frame_idx = 0\n",
        "\n",
        "    progress_bar = tqdm(total=num_frames)\n",
        "    progress_bar.set_description(\"Processing frames...\")\n",
        "\n",
        "    while cap.isOpened():\n",
        "        ret, frame = cap.read()\n",
        "        if not ret:\n",
        "            break\n",
        "        frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)\n",
        "        try:\n",
        "          pred_mask = pred_list[frame_idx]\n",
        "        except: break\n",
        "        masked_frame = draw_mask(frame,pred_mask)\n",
        "        # masked_frame = masked_pred_list[frame_idx]\n",
        "        masked_frame = cv2.cvtColor(masked_frame,cv2.COLOR_RGB2BGR)\n",
        "        out.write(masked_frame)\n",
        "        print('frame {} writed'.format(frame_idx),end='\\r')\n",
        "        frame_idx += 1\n",
        "        progress_bar.update(1)\n",
        "    out.release()\n",
        "    cap.release()\n",
        "    print(\"\\n{} saved\".format(io_args['output_video']))\n",
        "    print('\\nfinished')\n",
        "\n",
        "  if  save_gif:\n",
        "    # save colorized masks as a gif\n",
        "    imageio.mimsave(io_args['output_gif'],pred_list,fps=fps)\n",
        "    print(\"{} saved\".format(io_args['output_gif']))\n",
        "\n",
        "  from multiprocessing.pool import ThreadPool as Pool\n",
        "  from functools import partial\n",
        "  import PIL\n",
        "\n",
        "  threads = 12\n",
        "\n",
        "  def write_masks_frame(frame_num,  predicted_masks, output_folder, max_ids=255):\n",
        "    predicted_masks_frame = predicted_masks[frame_num]\n",
        "    for i in range(max_ids+1):\n",
        "      img_out = PIL.Image.fromarray(((predicted_masks_frame==i)*255).astype('uint8'))\n",
        "      img_out.save(os.path.join(output_folder, f'mask{i:03}', f'alpha_{frame_num:06}.jpg'))\n",
        "\n",
        "  def write_masks_frame_multi(predicted_masks, output_folder, max_ids):\n",
        "    for i in range(max_ids+1):\n",
        "      os.makedirs(os.path.join(output_folder, f'mask{i:03}'), exist_ok=True)\n",
        "\n",
        "    with Pool(threads) as p:\n",
        "      fn = partial(write_masks_frame, predicted_masks=predicted_masks, output_folder=output_folder, max_ids=max_ids)\n",
        "      result = list(tqdm(p.imap(fn, range(len(predicted_masks))), total=len(predicted_masks)))\n",
        "\n",
        "  if save_separate_masks:\n",
        "    print('Saving Separate masks')\n",
        "    write_masks_frame_multi(predicted_masks=pred_list, output_folder=output_multimask_dir, max_ids=segtracker.get_obj_num())\n",
        "    print(f'Saved masks to {output_multimask_dir}')\n",
        "\n"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 189,
          "referenced_widgets": [
            "a9d0125c361a426e9d5f8ac311f061c0",
            "a9db6b75095b4c979f23f96634cf7992",
            "24df4c73fa064e0f969ae40552335fe7",
            "969bf22f1b9e46bc98ae9fe7fce65086",
            "9dfc41ec4a984bbb9820dd6f07003535",
            "3aeeda7db5f8493b903ff06453156670",
            "0a4e220d8cdc44b4b62fdac6ca14e30c",
            "b109b007ea2d4fe29f3b4038961080a9",
            "22a376ea2da845db978aa6452e5a0c16",
            "1789f4e986da4dcc882345a7360fccb4",
            "7ba4c479c09b40e8bfbe9d68de309d05",
            "b8246408a9bc4d519b738cd6af2db010",
            "04e9b3573c374f76be694e4f4a326b95",
            "039a6f2543f044028cfd6d3dafcb1696",
            "5bbb7c1cd66a45a59870187f477bb6cb",
            "3a51c99a82174dab990205d6c0eec4b2",
            "f1e66c47264f4c31ba5a1afe74f3e7b2",
            "bf418989bcce4b16a8aaa8d9568a3cb1",
            "308b96514d91468fa86746da02629b74",
            "be5a771d7a774e8aa3448c3e474bec76",
            "47794863dd2a4280a9cdd460e2415df5",
            "1f1f7564317841cb8fe5e80b36cd61bd"
          ]
        },
        "cellView": "form",
        "id": "6nTX3C-RNvT5",
        "outputId": "d4f2179e-c914-49e3-a2dc-3e1ac78c4214"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "final text_encoder_type: bert-base-uncased\n",
            "Model loaded from ./ckpt/groundingdino_swint_ogc.pth \n",
            " => _IncompatibleKeys(missing_keys=[], unexpected_keys=['label_enc.weight', 'bert.embeddings.position_ids'])\n",
            "SegTracker has been initialized\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/51 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "a9d0125c361a426e9d5f8ac311f061c0"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saving Separate masks\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "text/plain": [
              "  0%|          | 0/51 [00:00<?, ?it/s]"
            ],
            "application/vnd.jupyter.widget-view+json": {
              "version_major": 2,
              "version_minor": 0,
              "model_id": "b8246408a9bc4d519b738cd6af2db010"
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Saved masks to /content/videoFrames/4a2baa1a78_masks\n"
          ]
        }
      ]
    }
  ]
}