demo/build_tsgm_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a5b28c59",
   "metadata": {},
   "source": [
    "# Demonstration for Building TSGM\n",
    "Note that this file uses Allensville environment in **gibson_tiny**  and you should prepare the dataset (glb, navmesh) in your habitat-lab scene_datasets path for running this demonstration.<br>\n",
    "check that the project directory indicates the **correct project path** (.../TopologicalSemanticGraphMemory)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a6842a8",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import os\n",
    "project_dir = \"/home/blackfoot/codes/TopologicalSemanticGraphMemory/demo\" #os.getcwd()\n",
    "project_dir = \"/\".join(project_dir.split(\"/\")[:-1])\n",
    "os.chdir(project_dir)\n",
    "project_dir"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9bd6886d",
   "metadata": {},
   "source": [
    "For visualization, please set the **display port** with your own. \n",
    "* os.environ['DISPLAY'] = 'localhost:10.0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dce4a742",
   "metadata": {},
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Set the display port!!\n",
    "\"\"\"\n",
    "os.environ['DISPLAY'] = 'localhost:10.0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b590771",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import torch\n",
    "import cv2, glob, imageio, joblib\n",
    "from PIL import Image\n",
    "from tqdm import tqdm \n",
    "import quaternion as q\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import habitat, habitat_sim\n",
    "habitat_path = habitat.__path__[0]\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "import torchvision.transforms as transforms\n",
    "import torch.nn as nn\n",
    "from utils.habitat_settings import default_sim_settings, make_cfg\n",
    "from utils.statics import CATEGORIES, COI_INDEX\n",
    "from matplotlib import animation\n",
    "from IPython.display import HTML\n",
    "from scipy.ndimage.interpolation import rotate\n",
    "from habitat.utils.visualizations import utils, fog_of_war\n",
    "normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],\n",
    "                                 std=[0.229, 0.224, 0.225])\n",
    "eval_augmentation = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    normalize\n",
    "])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b420c8b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.models import resnet18 as resnet18_img\n",
    "img_encoder = resnet18_img(num_classes=512)\n",
    "dim_mlp = img_encoder.fc.weight.shape[1]\n",
    "img_encoder.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), img_encoder.fc)\n",
    "ckpt_pth = os.path.join(project_dir, 'data/graph', 'Img_encoder.pth.tar')\n",
    "ckpt = torch.load(ckpt_pth, map_location='cpu')\n",
    "state_dict = {k[len('module.encoder_q.'):]: v for k, v in ckpt['state_dict'].items() if 'module.encoder_q.' in k}\n",
    "img_encoder.load_state_dict(state_dict)\n",
    "img_encoder.eval().cuda()\n",
    "def get_image_feat(rgb, depth, num_of_camera=12):\n",
    "    img_tensor = (torch.tensor(rgb[None]).cuda().float() / 255.).permute(0, 3, 1, 2)\n",
    "    feat = img_encoder(img_tensor)\n",
    "    vis_embedding = nn.functional.normalize(feat, dim=1)\n",
    "    return vis_embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f541aa52",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.Graph.resnet_obj import resnet18 as resnet18_obj\n",
    "feature_dim = 32\n",
    "object_encoder = resnet18_obj(num_classes=feature_dim)\n",
    "dim_mlp = object_encoder.fc.weight.shape[1]\n",
    "object_encoder.fc = nn.Sequential(nn.Linear(dim_mlp, dim_mlp), nn.ReLU(), object_encoder.fc)\n",
    "ckpt_pth = os.path.join(project_dir, 'data/graph', 'Obj_encoder.pth.tar')\n",
    "ckpt = torch.load(ckpt_pth, map_location='cpu')\n",
    "state_dict = {k[len('module.encoder_k.'):]: v for k, v in ckpt['state_dict'].items() if 'module.encoder_k.' in k}\n",
    "object_encoder.load_state_dict(state_dict)\n",
    "object_encoder.eval().cuda()\n",
    "def get_object_feat(rgb, bboxes, num_of_camera=12):\n",
    "    image = eval_augmentation(Image.fromarray(rgb.astype(np.uint8)))[None].cuda().float()\n",
    "    feat = object_encoder(image, torch.tensor(np.concatenate((np.zeros([len(bboxes), 1]), bboxes), -1)).cuda()[None].float())[0]\n",
    "    feat = nn.functional.normalize(feat, dim=1)\n",
    "    return feat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11ca07d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "settings = default_sim_settings.copy()\n",
    "settings[\"max_frames\"] = 100\n",
    "settings[\"width\"] = 256\n",
    "settings[\"height\"] = 256\n",
    "settings[\"scene\"] = ''\n",
    "settings[\"save_png\"] = False  # args.save_png\n",
    "settings[\"sensor_height\"] = 0.88\n",
    "settings[\"color_sensor\"] = True\n",
    "settings[\"semantic_sensor\"] = True\n",
    "settings[\"depth_sensor\"] = True\n",
    "settings[\"print_semantic_scene\"] = False\n",
    "settings[\"print_semantic_mask_stats\"] = False\n",
    "settings[\"compute_shortest_path\"] = False\n",
    "settings[\"compute_action_shortest_path\"] = False\n",
    "settings[\"panoramic_sensor\"] = True\n",
    "settings[\"seed\"] = 2343\n",
    "settings[\"silent\"] = False\n",
    "settings[\"enable_physics\"] = True\n",
    "settings[\"draw_lidar\"] = False\n",
    "settings[\"agent_radius\"] = 0.1\n",
    "settings[\"agent_height\"] = 1.2\n",
    "settings[\"multiview\"] = False\n",
    "settings[\"hfov\"] = 90\n",
    "settings[\"FORWARD_STEP_SIZE\"] = 0.25\n",
    "settings[\"TURN_ANGLE\"] = 10\n",
    "settings[\"tdv_height\"] = 512\n",
    "settings[\"tdv_width\"] = 512\n",
    "settings[\"allow_sliding\"] = True\n",
    "num_of_camera = 12\n",
    "pano_img_height = settings[\"height\"]//2\n",
    "img_height = pano_img_height\n",
    "img_width = float(img_height * 4)\n",
    "cam_width = float(pano_img_height * 4 // num_of_camera)\n",
    "settings[\"cam_width\"] = cam_width\n",
    "settings[\"pano_height\"] = pano_img_height\n",
    "\n",
    "dataset = 'gibson_tiny'\n",
    "scan_name = \"Allensville\"\n",
    "if dataset == \"mp3d\":\n",
    "    settings[\"scene\"] = os.path.join(habitat_path, '../data/scene_datasets/{}/{}/{}.glb'.format(dataset, scan_name, scan_name))\n",
    "elif \"gibson\" in dataset:\n",
    "    settings[\"scene\"] = os.path.join(habitat_path, '../data/scene_datasets/{}/{}.glb'.format(dataset, scan_name))\n",
    "elif dataset == \"hm3d\":\n",
    "    path = glob.glob(os.path.join(habitat_path, '../data/scene_datasets/{}/*/{}/{}.glb'.format(dataset, \"*\" + scan_name, scan_name)))[0]\n",
    "    settings[\"scene\"] = path\n",
    "cfg = make_cfg(settings)\n",
    "\n",
    "try:\n",
    "    sim.close()\n",
    "except:\n",
    "    pass\n",
    "\n",
    "sim = habitat_sim.Simulator(cfg)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9e52150",
   "metadata": {},
   "outputs": [],
   "source": [
    "max_object = 10\n",
    "bounds = sim.pathfinder.get_bounds()\n",
    "scene_objects = sim.semantic_scene.objects\n",
    "if dataset == \"mp3d\":\n",
    "    mapping = {int(obj.id.split(\"_\")[-1]): obj.category.index() for obj in scene_objects if obj != None}\n",
    "elif dataset == \"hm3d\":\n",
    "    mapping = {int(obj.id.split(\"_\")[-1]): obj.category.index() for obj in scene_objects if obj != None}\n",
    "else:\n",
    "    mapping = {int(obj.id.split(\"_\")[-1]): int(np.where([obj.category.name() == cat for cat in CATEGORIES['gibson']])[0][0]) for obj in scene_objects if obj != None}\n",
    "\n",
    "set_floor = 0\n",
    "floor = -1\n",
    "num_try = 0\n",
    "render_configs = joblib.load(os.path.join(project_dir, f\"data/floorplans/{dataset}_floorplans/render_config.pkl\"))\n",
    "while floor != set_floor:\n",
    "    init_position = sim.pathfinder.get_random_navigable_point()\n",
    "    floor =  int(np.argmin([abs(float(render_configs[scan_name][i]['z_low']) - init_position[1]) for i in render_configs[scan_name].keys()]))\n",
    "    num_try += 1\n",
    "    if num_try > 100:\n",
    "        print(\"Cannot find a valid position\")\n",
    "        break\n",
    "init_rotation = sim.get_agent(0).get_state().rotation.components\n",
    "scene_file = settings[\"scene\"]\n",
    "\n",
    "sim.reset()\n",
    "agent = sim.initialize_agent(0)\n",
    "start_state = agent.get_state()\n",
    "if (start_state.position != init_position).any():\n",
    "    start_state.position = init_position\n",
    "    start_state.rotation = q.from_float_array(init_rotation)  # self.init_rotation #\n",
    "    start_state.sensor_states = dict()  ## Initialize sensori\n",
    "agent.set_state(start_state)\n",
    "prev_position = agent.get_state().position\n",
    "prev_rotation = q.as_euler_angles(agent.state.rotation)[1]\n",
    "\n",
    "P = render_configs[scan_name][floor]['Projection']\n",
    "lower_bound, upper_bound = sim.pathfinder.get_bounds()\n",
    "imgWidth = round(float(render_configs[scan_name][floor]['width']))\n",
    "imgHeight = round(float(render_configs[scan_name][floor]['height']))\n",
    "world_min_width = float(render_configs[scan_name][floor]['x_low'])\n",
    "world_max_width = float(render_configs[scan_name][floor]['x_high'])\n",
    "world_min_height = float(render_configs[scan_name][floor]['y_low'])\n",
    "world_max_height = float(render_configs[scan_name][floor]['y_high'])\n",
    "worldWidth = abs(world_min_width) + abs(world_max_width)\n",
    "worldHeight = abs(world_min_height) + abs(world_max_height)\n",
    "\n",
    "print('====================================')\n",
    "print('scan_name: ', scan_name)\n",
    "print('floor: ', floor)\n",
    "print('lower_bound: ', lower_bound)\n",
    "print('upper_bound: ', upper_bound)\n",
    "print('imgWidth: ', imgWidth)\n",
    "print('imgHeight: ', imgHeight)\n",
    "\n",
    "ortho_map = cv2.imread(os.path.join(project_dir, f\"data/floorplans/{dataset}_floorplans/rgb/{scan_name}_level_{floor}.png\"))[...,::-1][...,:3]\n",
    "ortho_mask = cv2.imread(os.path.join(project_dir, f\"data/floorplans/{dataset}_floorplans/mask/{scan_name}_level_{floor}.png\"), 0)\n",
    "ortho_map[ortho_mask==0] = 255\n",
    "plt.axis(\"off\")\n",
    "plt.imshow(ortho_map)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a942b263",
   "metadata": {},
   "outputs": [],
   "source": [
    "scene_objects = sim.semantic_scene.objects\n",
    "object_loc = {int(obj.id.split(\"_\")[-1]): obj.aabb.center for obj in scene_objects if obj != None}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b4427ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImgGraph(object):\n",
    "    def __init__(self):\n",
    "        self.feature_dim = 512\n",
    "        self.M = 100\n",
    "\n",
    "    def num_node(self):\n",
    "        return len(self.node_position_list)\n",
    "\n",
    "    def reset(self):\n",
    "        self.node_position_list = []  # This position list is only for visualizations\n",
    "        self.node_rotation_list = []  # This rotation list is only for visualizations\n",
    "        self.graph_memory = np.zeros([self.M, self.feature_dim])\n",
    "        self.A = np.zeros([self.M, self.M], dtype=np.bool)\n",
    "        self.graph_mask = np.zeros(self.M)\n",
    "        self.graph_time = np.zeros(self.M)\n",
    "        self.last_localized_node_idx = np.zeros([1], dtype=np.int32)\n",
    "        self.last_localized_node_embedding = np.zeros([self.feature_dim], dtype=np.float32)\n",
    "\n",
    "    def initialize_graph(self, new_embeddings, positions, rotations):\n",
    "        self.add_node(node_idx=0, embedding=new_embeddings, time_step=0, position=positions, rotation=rotations)\n",
    "        self.record_localized_state(node_idx=0, embedding=new_embeddings)\n",
    "\n",
    "    def add_node(self, node_idx, embedding, time_step, position, rotation):\n",
    "        self.node_position_list.append(position)\n",
    "        self.node_rotation_list.append(rotation)\n",
    "        self.graph_memory[node_idx] = embedding\n",
    "        self.graph_mask[node_idx] = 1.0\n",
    "        self.graph_time[node_idx] = time_step\n",
    "\n",
    "    def record_localized_state(self, node_idx, embedding):\n",
    "        self.last_localized_node_idx = node_idx\n",
    "        self.last_localized_node_embedding = embedding\n",
    "\n",
    "    def add_edge(self, node_idx_a, node_idx_b):\n",
    "        self.A[node_idx_a, node_idx_b] = True\n",
    "        self.A[node_idx_b, node_idx_a] = True\n",
    "\n",
    "    def update_node(self, node_idx, time_info, embedding=None):\n",
    "        if embedding is not None:\n",
    "            self.graph_memory[node_idx] = embedding\n",
    "        self.graph_time[node_idx] = time_info\n",
    "\n",
    "    def get_positions(self, a=None):\n",
    "        if a is None:\n",
    "            return self.node_position_list\n",
    "        else:\n",
    "            return self.node_position_list[a]\n",
    "\n",
    "    def get_neighbor(self, node_idx, return_mask=False):\n",
    "        if return_mask:\n",
    "            return self.A[node_idx]\n",
    "        else:\n",
    "            return np.where(self.A[node_idx])[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9276152f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ObjGraph(object):\n",
    "    def __init__(self,):\n",
    "        self.feature_dim = 32\n",
    "        self.M = 1000\n",
    "        self.MV = 100\n",
    "\n",
    "    def num_node(self):\n",
    "        return len(self.node_position_list)\n",
    "\n",
    "    def reset(self):\n",
    "        self.node_position_list = []  # This position list is only for visualizations\n",
    "        self.graph_memory = np.zeros([self.M, self.feature_dim])\n",
    "        self.graph_category = np.zeros([self.M])\n",
    "        self.graph_score = np.zeros([self.M])\n",
    "        self.A_OV = np.zeros([self.M, self.MV], dtype=np.bool)\n",
    "        self.graph_mask = np.zeros(self.M)\n",
    "        self.graph_time = np.zeros([self.M], dtype=np.int32)\n",
    "        self.last_localized_node_idx = 0\n",
    "\n",
    "    def initialize_graph(self, new_embeddings, object_scores, object_categories, masks, positions):\n",
    "        if sum(masks == 1) == 0:\n",
    "            masks[0] = 1\n",
    "        self.add_node(node_idx=0, embedding=new_embeddings, object_score=object_scores, object_category=object_categories, time_step=0, mask=masks, position=positions, vis_node_idx=0)\n",
    "\n",
    "    def add_node(self, node_idx, embedding, object_score, object_category, mask, time_step, position, vis_node_idx):\n",
    "        node_idx_ = node_idx\n",
    "        i = 0\n",
    "        while True:\n",
    "            if mask[i] == 1:\n",
    "                self.node_position_list.append(position[i])\n",
    "                self.graph_memory[node_idx_] = embedding[i]\n",
    "                self.graph_score[node_idx_] = object_score[i]\n",
    "                self.graph_category[node_idx_] = object_category[i]\n",
    "                self.graph_mask[node_idx_] = 1.0\n",
    "                self.graph_time[node_idx_] = time_step\n",
    "                self.add_vo_edge([node_idx_], vis_node_idx)\n",
    "                node_idx_ += 1\n",
    "            i += 1\n",
    "            if i == len(position):\n",
    "                break\n",
    "\n",
    "    def add_vo_edge(self, node_idx_obj, curr_vis_node_idx):\n",
    "        for node_idx_obj_i in node_idx_obj:\n",
    "            self.A_OV[node_idx_obj_i, curr_vis_node_idx] = True\n",
    "\n",
    "    def update_node(self, node_idx, time_info, node_score, node_category, curr_vis_node_idx, embedding=None):\n",
    "        if embedding is not None:\n",
    "            self.graph_memory[node_idx] = embedding\n",
    "        self.graph_score[node_idx] = node_score\n",
    "        self.graph_category[node_idx] = node_category\n",
    "        self.graph_time[node_idx] = time_info\n",
    "        self.A_OV[node_idx, :] = False\n",
    "        self.A_OV[node_idx, curr_vis_node_idx] = True\n",
    "\n",
    "    def get_positions(self, a=None):\n",
    "        if a is None:\n",
    "            return self.node_position_list\n",
    "        else:\n",
    "            return self.node_position_list[a]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b29698b",
   "metadata": {},
   "outputs": [],
   "source": [
    "def is_close(embed_a, embed_b, return_prob=False, th=0.75):\n",
    "    logits = np.matmul(embed_a, embed_b.transpose(1, 0))\n",
    "    close = (logits > th)\n",
    "    if return_prob:\n",
    "        return close, logits\n",
    "    else:\n",
    "        return close"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfd39007",
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_image_graph(imggraph, objgraph, new_embedding, curr_obj_embeding, \n",
    "                       object_score, object_category, position, rotation, time, done, img_node_th=0.8, obj_node_th=0.8):\n",
    "    # The position is only used for visualizations\n",
    "    if done:\n",
    "        imggraph.reset()\n",
    "        imggraph.initialize_graph(new_embedding, position, rotation)\n",
    "\n",
    "    obj_close = True\n",
    "    obj_graph_mask = objgraph.graph_score[objgraph.A_OV[:, imggraph.last_localized_node_idx]] > 0.5\n",
    "    if len(obj_graph_mask) > 0:\n",
    "        curr_obj_mask = object_score > 0.5\n",
    "        if np.sum(curr_obj_mask) / len(curr_obj_mask) >= 0.5:\n",
    "            close_obj, prob_obj = is_close(objgraph.graph_memory[objgraph.A_OV[:, imggraph.last_localized_node_idx]], curr_obj_embeding, return_prob=True, th=obj_node_th)\n",
    "            close_obj = close_obj[obj_graph_mask, :][:, curr_obj_mask]\n",
    "            category_mask = objgraph.graph_category[objgraph.A_OV[:, imggraph.last_localized_node_idx]][obj_graph_mask][:, None] == object_category[curr_obj_mask]\n",
    "            close_obj[~category_mask] = False\n",
    "            if len(close_obj) >= 3:\n",
    "                clos_obj_p = close_obj.any(1).sum() / (close_obj.shape[0])\n",
    "                if clos_obj_p < 0.1:  # Fail to localize (find the same object) with the last localized frame\n",
    "                    obj_close = False\n",
    "\n",
    "    close, prob = is_close(imggraph.last_localized_node_embedding[None], new_embedding[None], return_prob=True, th=img_node_th)\n",
    "\n",
    "    found = (np.array(done) + close.squeeze()) & np.array(obj_close).squeeze()\n",
    "    found_prev = False\n",
    "    found = found\n",
    "    found_in_memory = False\n",
    "    to_add = False\n",
    "    if found:\n",
    "        imggraph.update_node(imggraph.last_localized_node_idx, time)\n",
    "        found_prev = True\n",
    "    else:\n",
    "        check_list = 1 - imggraph.graph_mask[:imggraph.num_node()]\n",
    "        check_list[imggraph.last_localized_node_idx] = 1.0\n",
    "        while not found:\n",
    "            not_checked_yet = np.where((1 - check_list))[0]\n",
    "            neighbor_embedding = imggraph.graph_memory[not_checked_yet]\n",
    "            num_to_check = len(not_checked_yet)\n",
    "            if num_to_check == 0:\n",
    "                to_add = True\n",
    "                break\n",
    "            else:\n",
    "                close, prob = is_close(new_embedding[None], neighbor_embedding, return_prob=True, th=img_node_th)\n",
    "                close = close[0];\n",
    "                prob = prob[0]\n",
    "                close_idx = np.where(close)[0]\n",
    "                if len(close_idx) >= 1:\n",
    "                    found_node = not_checked_yet[prob.argmax()]\n",
    "                else:\n",
    "                    found_node = None\n",
    "                if found_node is not None:\n",
    "                    found = True\n",
    "                    if abs(time - imggraph.graph_time[found_node]) > 20:\n",
    "                        found_in_memory = True\n",
    "                    imggraph.update_node(found_node, time, new_embedding)\n",
    "                    imggraph.add_edge(found_node, imggraph.last_localized_node_idx)\n",
    "                    imggraph.record_localized_state(found_node, new_embedding)\n",
    "                check_list[found_node] = 1.0\n",
    "\n",
    "    if to_add:\n",
    "        new_node_idx = imggraph.num_node()\n",
    "        imggraph.add_node(new_node_idx, new_embedding, time, position, rotation)\n",
    "        imggraph.add_edge(new_node_idx, imggraph.last_localized_node_idx)\n",
    "        imggraph.record_localized_state(new_node_idx, new_embedding)\n",
    "    last_localized_node_idx = imggraph.last_localized_node_idx\n",
    "    return imggraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f75589b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "def update_object_graph(imggraph, objgraph, object_embedding, object_score, object_category, object_mask, object_position, time, done, obj_node_th=0.8):\n",
    "    # The position is only used for visualizations. Remove if object features are similar\n",
    "    # object masking\n",
    "    object_score = object_score[object_mask == 1]\n",
    "    object_category = object_category[object_mask == 1]\n",
    "    object_position = object_position[object_mask == 1]\n",
    "    object_embedding = object_embedding[object_mask == 1]\n",
    "    object_mask = object_mask[object_mask == 1]\n",
    "    if done:\n",
    "        objgraph.reset()\n",
    "        objgraph.initialize_graph(object_embedding, object_score, object_category, object_mask, object_position)\n",
    "\n",
    "    to_add = [True] * int(sum(object_mask))\n",
    "    not_found = not done  # Dense connection mode\n",
    "    if not_found:\n",
    "        hop1_vis_node = imggraph.A[imggraph.last_localized_node_idx]\n",
    "        hop1_obj_node_mask = np.sum(objgraph.A_OV.transpose(1, 0)[hop1_vis_node], 0) > 0\n",
    "        curr_obj_node_mask = objgraph.A_OV[:, imggraph.last_localized_node_idx]\n",
    "        neighbor_obj_node_mask = (hop1_obj_node_mask + curr_obj_node_mask) > 0\n",
    "        neighbor_node_embedding = objgraph.graph_memory[neighbor_obj_node_mask]\n",
    "        neighbor_obj_memory_idx = np.where(neighbor_obj_node_mask)[0]\n",
    "        neighbor_obj_memory_score = objgraph.graph_score[neighbor_obj_memory_idx]\n",
    "        neighbor_obj_memory_cat = objgraph.graph_category[neighbor_obj_memory_idx]\n",
    "\n",
    "        close, prob = is_close(neighbor_node_embedding, object_embedding, return_prob=True, th=obj_node_th)\n",
    "        for c_i in range(prob.shape[1]):\n",
    "            close_mem_indices = np.where(close[:, c_i] == 1)[0]\n",
    "            for m_i in close_mem_indices:\n",
    "                is_same = False\n",
    "                to_update = False\n",
    "                if (object_category[c_i] == neighbor_obj_memory_cat[m_i]) and object_category[c_i] != -1:\n",
    "                    is_same = True\n",
    "                    if object_score[c_i] > neighbor_obj_memory_score[m_i]:\n",
    "                        to_update = True\n",
    "\n",
    "                if is_same:\n",
    "                    to_add[c_i] = False\n",
    "\n",
    "                if to_update:\n",
    "                    objgraph.update_node(m_i, time, object_score[c_i], object_category[c_i], int(imggraph.last_localized_node_idx), object_embedding[c_i])\n",
    "                    break\n",
    "\n",
    "        # Add new objects to graph\n",
    "        if sum(to_add) > 0:\n",
    "            start_node_idx = objgraph.num_node()\n",
    "            new_idx = np.where(np.stack(to_add))[0]\n",
    "            objgraph.add_node(start_node_idx, object_embedding[new_idx], object_score[new_idx],\n",
    "                              object_category[new_idx], object_mask[new_idx], time,\n",
    "                              object_position[new_idx], int(imggraph.last_localized_node_idx))\n",
    "    return objgraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de860b83",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_graph_on_map(topdownmap, imggraph, objgraph, draw_im_graph=True, draw_obj_graph=True, node_size_px=60):\n",
    "    node_list = imggraph.node_position_list\n",
    "    affinity = imggraph.A\n",
    "    last_localized_imnode = imggraph.last_localized_node_idx\n",
    "    print(last_localized_imnode)\n",
    "    obj_node_list = objgraph.node_position_list\n",
    "    ov_affinity = objgraph.A_OV\n",
    "    if draw_obj_graph:\n",
    "        h, w, _ = topdownmap.shape\n",
    "        draw_obj_point_list = []\n",
    "        for idx, node_position in enumerate(obj_node_list):\n",
    "            connected_imnode = np.where(ov_affinity[idx])[0]\n",
    "            draw_obj_point_list.append(node_position)\n",
    "\n",
    "            if draw_im_graph:\n",
    "                neighbors = np.where(ov_affinity[idx])[0]\n",
    "                for neighbor_idx in neighbors:\n",
    "                    neighbor_position = node_list[neighbor_idx]\n",
    "                    node_map_pose = to_grid(node_position[0], node_position[2])\n",
    "                    neighbor_map_pose = to_grid(neighbor_position[0], neighbor_position[2])\n",
    "                    line = cv2.line(\n",
    "                        topdownmap,\n",
    "                        node_map_pose,\n",
    "                        neighbor_map_pose,\n",
    "                        [253, 173, 211],\n",
    "                        thickness=10\n",
    "                    )\n",
    "                    alpha = 0.8\n",
    "                    topdownmap = cv2.addWeighted(line, alpha, topdownmap, 1 - alpha, 0)\n",
    "\n",
    "    if draw_im_graph:\n",
    "        draw_im_point_list = []\n",
    "        for idx, node_position in enumerate(node_list):\n",
    "            if idx == last_localized_imnode:\n",
    "                draw_im_point_list.append([node_position, [51, 153, 102]])\n",
    "            else:\n",
    "                draw_im_point_list.append([node_position, [15, 119, 143]])\n",
    "            neighbors = np.where(affinity[idx])[0]\n",
    "            for neighbor_idx in neighbors:\n",
    "                neighbor_position = node_list[neighbor_idx]\n",
    "                node_map_pose = to_grid(node_position[0], node_position[2])\n",
    "                neighbor_map_pose = to_grid(neighbor_position[0], neighbor_position[2])\n",
    "                cv2.line(\n",
    "                    topdownmap,\n",
    "                    node_map_pose,\n",
    "                    neighbor_map_pose,\n",
    "                    [177, 232, 246],\n",
    "                    thickness=30,\n",
    "                )\n",
    "\n",
    "    if draw_im_graph:\n",
    "        for node_position, node_color in draw_im_point_list:\n",
    "            graph_node_center = to_grid(node_position[0], node_position[2])\n",
    "            cv2.circle(topdownmap, graph_node_center, node_size_px // 3, node_color, -1)\n",
    "\n",
    "    if draw_obj_graph:\n",
    "        cnt = 0\n",
    "        for node_position in draw_obj_point_list:\n",
    "            graph_node_center = to_grid(node_position[0], node_position[2])\n",
    "            triangle_cnt = np.array(graph_node_center + np.array([pt1, pt2, pt3]).astype(np.int32))\n",
    "            cv2.drawContours(topdownmap, [triangle_cnt], 0, [248, 106, 176], -1)\n",
    "            cnt += 1\n",
    "\n",
    "    return topdownmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4baca37",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_objects(semantic, mapping, depth, img_width, img_height):\n",
    "    semantic = semantic.astype(np.int32)\n",
    "    max_key = np.max(np.array(list(mapping.keys())))\n",
    "    replace_values = []\n",
    "    for i in np.arange(max_key + 1):\n",
    "        try:\n",
    "            replace_values.append(mapping[i])\n",
    "        except:\n",
    "            replace_values.append(-1)\n",
    "    semantic_obs_class = np.take(replace_values, semantic)\n",
    "    COI_MASK = [(semantic_obs_class == ci).astype(np.int32) for ci in COI_INDEX[dataset.split(\"_\")[0]]]  # class mask\n",
    "    unique_instances = np.unique(semantic * np.sum(np.stack(COI_MASK), 0))[1:]\n",
    "    bbox_ids = unique_instances\n",
    "    instance_segment = [(semantic == i).astype(np.int32) for i in unique_instances]\n",
    "    box_categories = [np.unique(semantic_obs_class[semantic == i])[0] for i in unique_instances]\n",
    "    if len(instance_segment) > 0:\n",
    "        object_size = np.stack([np.sum(instance_segman) for instance_segman in instance_segment])\n",
    "        mask = (object_size > 100)\n",
    "        instance_segment = [instance_segman for i, instance_segman in enumerate(instance_segment) if mask[i] == 1]\n",
    "        box_categories = np.stack(box_categories)[mask == 1]\n",
    "        bbox_ids = np.array(bbox_ids)[mask == 1]\n",
    "\n",
    "    x1s = [np.min(np.where(instance_segment[i])[1]) for i in range(len(instance_segment))]\n",
    "    y1s = [np.min(np.where(instance_segment[i])[0]) for i in range(len(instance_segment))]\n",
    "    x2s = [np.max(np.where(instance_segment[i])[1]) for i in range(len(instance_segment))]\n",
    "    y2s = [np.max(np.where(instance_segment[i])[0]) for i in range(len(instance_segment))]\n",
    "    bboxes = np.stack((x1s, y1s, x2s, y2s), 1)\n",
    "    if len(bboxes) > 0:\n",
    "        edge_box_idx = np.where(bboxes[:, 2] - bboxes[:, 0] > img_width * 0.8)[0]\n",
    "        not_edge_box_idx = np.where(bboxes[:, 2] - bboxes[:, 0] <= img_width * 0.8)[0]\n",
    "        if len(edge_box_idx) > 0:\n",
    "            x1s1 = [np.min(np.where(instance_segment[i][:, :int(instance_segment[i].shape[1] / 2)])[1]) for i in edge_box_idx]\n",
    "            y1s1 = [np.min(np.where(instance_segment[i][:, :int(instance_segment[i].shape[1] / 2)])[0]) for i in edge_box_idx]\n",
    "            x2s1 = [np.max(np.where(instance_segment[i][:, :int(instance_segment[i].shape[1] / 2)])[1]) for i in edge_box_idx]\n",
    "            y2s1 = [np.max(np.where(instance_segment[i][:, :int(instance_segment[i].shape[1] / 2)])[0]) for i in edge_box_idx]\n",
    "            bboxes_1 = np.stack((x1s1, y1s1, x2s1, y2s1), 1)\n",
    "            bboxes_1_categories = box_categories[edge_box_idx]\n",
    "            bboxes_1_ids = bbox_ids[edge_box_idx]\n",
    "            x1s2 = [int(instance_segment[i].shape[1] / 2) + np.min(np.where(instance_segment[i][:, int(instance_segment[i].shape[1] / 2):])[1]) for i in edge_box_idx]\n",
    "            y1s2 = [np.min(np.where(instance_segment[i][:, int(instance_segment[i].shape[1] / 2):])[0]) for i in edge_box_idx]\n",
    "            x2s2 = [int(instance_segment[i].shape[1] / 2) + np.max(np.where(instance_segment[i][:, int(instance_segment[i].shape[1] / 2):])[1]) for i in edge_box_idx]\n",
    "            y2s2 = [np.max(np.where(instance_segment[i][:, int(instance_segment[i].shape[1] / 2):])[0]) for i in edge_box_idx]\n",
    "            bboxes_2 = np.stack((x1s2, y1s2, x2s2, y2s2), 1)\n",
    "            bboxes_2_categories = box_categories[edge_box_idx]\n",
    "            bboxes_2_ids = bbox_ids[edge_box_idx]\n",
    "            bboxes_ = bboxes[not_edge_box_idx]\n",
    "            box_categories_ = box_categories[not_edge_box_idx]\n",
    "            bbox_ids_ = bbox_ids[not_edge_box_idx]\n",
    "            bboxes = np.concatenate((bboxes_, bboxes_1, bboxes_2), 0)\n",
    "            box_categories = np.concatenate((box_categories_, bboxes_1_categories, bboxes_2_categories), 0)\n",
    "            bbox_ids = np.concatenate((bbox_ids_, bboxes_1_ids, bboxes_2_ids), 0)\n",
    "\n",
    "    if len(depth) > 0 and len(bboxes) > 0:\n",
    "        box_pix_xs, box_pix_ys = ((bboxes[:, 0])).astype(np.int32), ((bboxes[:, 1])).astype(np.int32)\n",
    "        box_depth = np.stack([depth[box_pix_y, box_pix_x] for box_pix_x, box_pix_y in zip(box_pix_xs, box_pix_ys)])\n",
    "        box_mask = box_depth < 5\n",
    "        if np.sum(box_mask) == 0:\n",
    "            return [], [], [], [], []\n",
    "        bboxes = bboxes[box_mask]\n",
    "        box_categories = box_categories[box_mask]\n",
    "        bbox_ids = bbox_ids[box_mask]\n",
    "        box_depth = box_depth[box_mask]\n",
    "        box_world = np.stack([object_loc[bbox_id] for bbox_id in bbox_ids])\n",
    "\n",
    "        bboxes = bboxes.astype(np.float32)\n",
    "        bboxes[:, 0] = bboxes[:, 0].astype(np.float32) / float(img_width)\n",
    "        bboxes[:, 1] = bboxes[:, 1].astype(np.float32) / float(img_height)\n",
    "        bboxes[:, 2] = bboxes[:, 2].astype(np.float32) / float(img_width)\n",
    "        bboxes[:, 3] = bboxes[:, 3].astype(np.float32) / float(img_height)\n",
    "        return bboxes, box_categories, bbox_ids, box_depth, box_world\n",
    "    return [], [], [], [], []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77d933ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "size= 0.4\n",
    "pt1 = np.array([150-150, 100-150])*size\n",
    "pt2 = np.array([100-150, 200-150])*size\n",
    "pt3 = np.array([200-150, 200-150])*size\n",
    "pt1 = pt1.astype(np.int32)\n",
    "pt2 = pt2.astype(np.int32)\n",
    "pt2 = pt2.astype(np.int32)\n",
    "JACKAL_SPRITE = imageio.imread(os.path.join(project_dir, \"data/assets/maps_topdown_agent_sprite/jackal.png\"))\n",
    "JACKAL_SPRITE = np.ascontiguousarray(np.flipud(JACKAL_SPRITE))\n",
    "initial_agent_size = JACKAL_SPRITE.shape[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "444ae3d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_grid(realworld_x: float,realworld_y: float,):\n",
    "    A = [realworld_x-(worldWidth+2*world_min_width)/2, realworld_y-(worldHeight+2*world_min_height)/2, 1, 1]\n",
    "    grid_x, grid_y = np.array([imgWidth/2, imgHeight/2]) * np.matmul(P, A)[:2] + np.array([imgWidth/2, imgHeight/2])\n",
    "    return int(grid_x), int(grid_y)\n",
    "\n",
    "def from_grid(grid_x: int,grid_y: int):\n",
    "    realworld_x, realworld_y = np.matmul(np.linalg.inv(P), [(2 * grid_x - imgWidth)/imgWidth, (2 * grid_y - imgHeight)/imgHeight, 1, 1])[:2] \\\n",
    "                               + np.array([(worldWidth+2*world_min_width)/2, (worldHeight+2*world_min_height)/2])\n",
    "    return realworld_x, realworld_y\n",
    "from habitat.utils.geometry_utils import (\n",
    "    quaternion_from_coeff,\n",
    "    quaternion_rotate_vector,\n",
    ")\n",
    "from habitat.tasks.utils import cartesian_to_polar\n",
    "def get_map_angle(ref_rotation):\n",
    "    heading_vector = quaternion_rotate_vector(\n",
    "        ref_rotation.inverse(), np.array([0, 0, -1])\n",
    "    )\n",
    "\n",
    "    phi = cartesian_to_polar(-heading_vector[2], heading_vector[0])[1]\n",
    "    return np.array(phi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "761bdcbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_bbox(rgb: np.ndarray, bboxes: np.ndarray, color = (178,193,118), is_detection=False) -> np.ndarray:\n",
    "    for i, bbox in enumerate(bboxes):\n",
    "        imgHeight, imgWidth, _ = rgb.shape\n",
    "        cv2.rectangle(rgb, (int(bbox[0]*imgWidth), int(bbox[1]*imgHeight)), (int(bbox[2]*imgWidth), int(bbox[3]*imgHeight)), (255, 255, 0), 5)\n",
    "    return rgb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ea3e41f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def step(in_action):\n",
    "    action_names = list(cfg.agents[settings[\"default_agent\"]].action_space.keys())\n",
    "    action = action_names[in_action]\n",
    "    sim.step(action)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8de6962f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "action_names = list(cfg.agents[settings[\"default_agent\"]].action_space.keys())\n",
    "action_names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1c772ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from magnum import Quaternion, Vector3, Rad\n",
    "start_position = Vector3([from_grid(356, 535)[0], init_position[1], from_grid(356, 535)[1]])\n",
    "goal_position = Vector3([from_grid(187, 124)[0], init_position[1], from_grid(187, 124)[1]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "582ef0b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "agent = sim.initialize_agent(0)\n",
    "start_state = agent.get_state()\n",
    "start_state.position = start_position\n",
    "agent.set_state(start_state)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fae8b2ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def draw_agent(topdownmap, agent_position, agent_rotation, jackal_radius_px=50):\n",
    "    rotated_jackal = rotate(\n",
    "        JACKAL_SPRITE, agent_rotation * 180 / np.pi + 90\n",
    "    )\n",
    "    initial_jackal_size = JACKAL_SPRITE.shape[0]\n",
    "    new_size = rotated_jackal.shape[0]\n",
    "    jackal_size_px = max(1, int(jackal_radius_px * 2 * new_size / initial_agent_size))\n",
    "    resized_jackal = cv2.resize(\n",
    "        rotated_jackal,\n",
    "        (jackal_size_px, jackal_size_px),\n",
    "        interpolation=cv2.INTER_LINEAR\n",
    "    )\n",
    "    topdownmap = utils.paste_overlapping_image(topdownmap, resized_jackal, agent_position)\n",
    "    return topdownmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0d2bd61",
   "metadata": {},
   "outputs": [],
   "source": [
    "def concat_vertical(img1, img2):\n",
    "    img1_shape = img1.shape\n",
    "    img2_shape = img2.shape\n",
    "    if img1_shape[1] > img2_shape[1]:\n",
    "        img2 = cv2.resize(img2, (img1_shape[1], int(img2_shape[0] / img2_shape[1] * img1_shape[1])))\n",
    "    elif img1_shape[1] < img2_shape[1]:\n",
    "        img1 = cv2.resize(img1, (img2_shape[1], int(img1_shape[0] / img1_shape[1] * img2_shape[1])))\n",
    "    return np.concatenate((img1, img2), axis=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "069a19b2",
   "metadata": {},
   "source": [
    "# Demonstration\n",
    "Use 'w/a/d' for exploring the environment and press 'q' for stop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "id": "96ede3a9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "imggraph = ImgGraph()\n",
    "objgraph = ObjGraph()\n",
    "positions = []\n",
    "rotations = []\n",
    "topdownmap = ortho_map.copy()\n",
    "obs = sim.get_sensor_observations()\n",
    "rgb = np.concatenate([obs['panoramic_rgb_part_sensor_%d' % (i)][:, :, :3] for i in range(num_of_camera)], 1)\n",
    "semantic_array = np.concatenate([obs['panoramic_semantic_part_sensor_%d' % (i)] for i in range(num_of_camera)], 1)\n",
    "depth_array = np.concatenate([obs['panoramic_depth_part_sensor_%d' % (i)] for i in range(num_of_camera)], 1)\n",
    "bboxes, bbox_categories, bbox_ids, bbox_depth, bbox_pose = get_objects(semantic_array, mapping, depth_array, rgb.shape[1], rgb.shape[0])\n",
    "\n",
    "position = sim.get_agent(0).get_state().position\n",
    "rotation = sim.get_agent(0).get_state().rotation.components\n",
    "cur_map_coord = to_grid(position[0], position[2])[::-1]\n",
    "cur_map_rot = get_map_angle(sim.get_agent(0).get_state().rotation)\n",
    "image_feat = get_image_feat(rgb, depth_array).cpu().detach().numpy()\n",
    "object_feat = []\n",
    "if len(bboxes) > 0:\n",
    "    object_feat = get_object_feat(rgb, bboxes).cpu().detach().numpy()\n",
    "position = sim.get_agent(0).get_state().position\n",
    "rotation = sim.get_agent(0).get_state().rotation.components\n",
    "imggraph.reset()\n",
    "imggraph.initialize_graph(image_feat[0], position, rotation)\n",
    "objgraph.reset()\n",
    "if len(bboxes) > 0:\n",
    "    objgraph.initialize_graph(object_feat, np.ones(len(bboxes)), bbox_categories,np.ones(len(bboxes)), bbox_pose)\n",
    "topdownmap = draw_agent(topdownmap, cur_map_coord, cur_map_rot)\n",
    "\n",
    "cv2.imshow(f'vis', rgb[:, :, [2, 1, 0]])\n",
    "cv2.imshow(f'graph', topdownmap[::2, ::2, [2, 1, 0]])\n",
    "\n",
    "topdownmaps = []\n",
    "images = []\n",
    "t = 0\n",
    "while True:\n",
    "    topdownmap = ortho_map.copy()\n",
    "    key = cv2.waitKey(0)\n",
    "    if key == ord('w'):\n",
    "        step(1)\n",
    "    elif key == ord('a'):\n",
    "        step(2)\n",
    "    elif key == ord('d'):\n",
    "        step(3)\n",
    "    elif key == ord('q'):\n",
    "        cv2.destroyAllWindows()\n",
    "        plt.imshow(topdownmap[::2, ::2, [2, 1, 0]])\n",
    "        break\n",
    "    position = sim.get_agent(0).get_state().position\n",
    "    rotation = sim.get_agent(0).get_state().rotation.components\n",
    "    cur_map_coord = to_grid(position[0], position[2])[::-1]\n",
    "    cur_map_rot = get_map_angle(sim.get_agent(0).get_state().rotation)\n",
    "    positions.append(position)\n",
    "    rotations.append(rotation)\n",
    "    obs = sim.get_sensor_observations()\n",
    "    rgb = np.concatenate([obs['panoramic_rgb_part_sensor_%d' % (i)][:, :, :3] for i in range(num_of_camera)], 1)\n",
    "    semantic_array = np.concatenate([obs['panoramic_semantic_part_sensor_%d' % (i)] for i in range(num_of_camera)], 1)\n",
    "    depth_array = np.concatenate([obs['panoramic_depth_part_sensor_%d' % (i)] for i in range(num_of_camera)], 1) / 10.\n",
    "    bboxes, bbox_categories, bbox_ids, bbox_depth, bbox_pose = get_objects(semantic_array, mapping, depth_array, rgb.shape[1], rgb.shape[0])\n",
    "\n",
    "    image_feat = get_image_feat(rgb, depth_array).cpu().detach().numpy()\n",
    "    object_feat = []\n",
    "    if len(bboxes) > 0:\n",
    "        object_feat = get_object_feat(rgb, bboxes).cpu().detach().numpy()\n",
    "    \n",
    "    imggraph = update_image_graph(imggraph, objgraph, image_feat[0], object_feat, np.ones(len(bboxes)), bbox_categories, \n",
    "                                  position, rotation, time=t, done=False)\n",
    "    \n",
    "    if len(bboxes) > 0:\n",
    "        objgraph = update_object_graph(imggraph, objgraph, object_feat, np.ones(len(bboxes)), bbox_categories, np.ones(len(bboxes)), \n",
    "                                       bbox_pose, time=t, done=False)\n",
    "        \n",
    "    topdownmap = draw_graph_on_map(topdownmap, imggraph, objgraph)\n",
    "    topdownmap = draw_agent(topdownmap, cur_map_coord, cur_map_rot)\n",
    "    rgb = draw_bbox(rgb, bboxes)\n",
    "    topdownmaps.append(concat_vertical(rgb, topdownmap))\n",
    "    cv2.imshow(f'vis', rgb[:, :, [2, 1, 0]])\n",
    "    cv2.imshow(f'graph', topdownmap[::2, ::2, [2, 1, 0]])    \n",
    "    t = t + 1    \n",
    "topdownmaps = np.stack(topdownmaps)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16ef925f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def init():\n",
    "    im.set_data(topdownmaps[0])\n",
    "\n",
    "def animate(i):\n",
    "    im.set_data(topdownmaps[i])\n",
    "    return im"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d26c7a8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(10,10))\n",
    "plt.axis('off')\n",
    "plt.gca().set_axis_off()\n",
    "plt.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)\n",
    "plt.margins(0,0)\n",
    "\n",
    "im = plt.imshow(topdownmaps[0])\n",
    "plt.close() # this is required to not display the generated image\n",
    "anim = animation.FuncAnimation(fig, animate, init_func=init, frames=topdownmaps.shape[0], interval=50)\n",
    "HTML(anim.to_html5_video())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "277eab73",
   "metadata": {},
   "outputs": [],
   "source": [
    "writergif = animation.PillowWriter(fps=20) \n",
    "anim.save(os.path.join(project_dir, \"demo/tsgm_demo.gif\"), writer=writergif)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:habitat021]",
   "language": "python",
   "name": "conda-env-habitat021-py"
  },
  "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.7.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}