interactive_simplicial_complex_from_attention_on_images.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "14a2be8d-efa2-474f-a51d-d0cec141f400",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from transformers import ViTModel, ViTFeatureExtractor\n",
    "import gudhi as gd\n",
    "import networkx as nx\n",
    "from scipy.spatial.distance import jensenshannon\n",
    "import plotly.graph_objs as go\n",
    "from PIL import Image\n",
    "import ipywidgets as widgets\n",
    "\n",
    "\n",
    "def get_attention_matrix(image, model, tokenizer, layer, head):\n",
    "    inputs = tokenizer(images=[image], return_tensors=\"pt\")\n",
    "    outputs = model(**inputs, output_attentions=True)\n",
    "    attention = outputs.attentions[layer][0, head].detach().cpu().numpy()\n",
    "    return attention\n",
    "\n",
    "def compute_persistence(attention_matrix):\n",
    "    softmax_attention = np.exp(attention_matrix) / np.sum(np.exp(attention_matrix), axis=-1)[:, np.newaxis]\n",
    "    distance_matrix = np.array([[np.sqrt(jensenshannon(softmax_attention[i], softmax_attention[j])) for j in range(softmax_attention.shape[0])] for i in range(softmax_attention.shape[0])])\n",
    "    \n",
    "    rips_complex = gd.RipsComplex(distance_matrix=distance_matrix, max_edge_length=np.inf)\n",
    "    simplex_tree = rips_complex.create_simplex_tree(max_dimension=2)\n",
    "    persistence = simplex_tree.persistence(min_persistence=0.01)\n",
    "    return persistence, simplex_tree, distance_matrix\n",
    "\n",
    "def process_image(image_path, model, tokenizer, layer, head):\n",
    "    image = Image.open(image_path)\n",
    "    attention_matrix = get_attention_matrix(image, model, tokenizer, layer, head)\n",
    "    persistence, simplex_tree, distance_matrix = compute_persistence(attention_matrix)\n",
    "    tokens = list(range(attention_matrix.shape[0]))\n",
    "    return persistence, simplex_tree, distance_matrix, tokens\n",
    "\n",
    "def plot_simplicial_complex_3d(simplex_tree, distance_matrix, title, threshold, tokens):\n",
    "    g = nx.Graph()\n",
    "    for (simplex, _) in simplex_tree.get_filtration():\n",
    "        if len(simplex) == 2:\n",
    "            if distance_matrix[simplex[0]][simplex[1]] <= threshold:\n",
    "                g.add_edge(simplex[0], simplex[1])\n",
    "\n",
    "    labels = {node: tokens[node] for node in g.nodes()}\n",
    "    \n",
    "    pos = nx.spring_layout(g, dim=3, seed=42)\n",
    "    \n",
    "    Xn = [pos[k][0] for k in g.nodes()]\n",
    "    Yn = [pos[k][1] for k in g.nodes()]\n",
    "    Zn = [pos[k][2] for k in g.nodes()]\n",
    "    \n",
    "    Xe = []\n",
    "    Ye = []\n",
    "    Ze = []\n",
    "    for e in g.edges():\n",
    "        Xe += [pos[e[0]][0], pos[e[1]][0], None]\n",
    "        Ye += [pos[e[0]][1], pos[e[1]][1], None]\n",
    "        Ze += [pos[e[0]][2], pos[e[1]][2], None]\n",
    "    \n",
    "    trace_edges = go.Scatter3d(x=Xe, y=Ye, z=Ze, mode='lines', line=dict(color='gray', width=1))\n",
    "    trace_nodes = go.Scatter3d(x=Xn, y=Yn, z=Zn, mode='markers+text', text=list(labels.values()), marker=dict(symbol='circle', size=10, color='lightblue'), textposition=\"top center\")\n",
    "    layout = go.Layout(title=title, scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'), showlegend=False)\n",
    "    fig = go.Figure(data=[trace_edges, trace_nodes], layout=layout)\n",
    "    fig.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "65404cd8-03d6-45ea-9180-a4bcdc68a2f4",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at google/vit-base-patch16-224 were not used when initializing ViTModel: ['classifier.weight', 'classifier.bias']\n",
      "- This IS expected if you are initializing ViTModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing ViTModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
      "Some weights of ViTModel were not initialized from the model checkpoint at google/vit-base-patch16-224 and are newly initialized: ['vit.pooler.dense.weight', 'vit.pooler.dense.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "0344ed0fbd4841a3a4c5ddf5b1d8aae7",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.001, continuous_update=False, description='Threshold:', max=0.2, min…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.update_plot(threshold)>"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load pre-trained visual transformer model and tokenizer\n",
    "tokenizer = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224')\n",
    "model = ViTModel.from_pretrained('google/vit-base-patch16-224')\n",
    "\n",
    "# Image paths\n",
    "image_paths = [\"heart.jpg\", \"platonic_metatron.jpg\"]\n",
    "\n",
    "# Choose a layer and head\n",
    "layer = 1\n",
    "head = 2\n",
    "\n",
    "# Process images\n",
    "results = [process_image(image_path, model, tokenizer, layer, head) for image_path in image_paths]\n",
    "\n",
    "# Create a slider for the threshold value\n",
    "threshold_slider = widgets.FloatSlider(\n",
    "    value=0.001,\n",
    "    min=0.00001,\n",
    "    max=0.2,\n",
    "    step=0.00001,\n",
    "    description='Threshold:',\n",
    "    continuous_update=False\n",
    ")\n",
    "\n",
    "# Define a function to update the plot based on the threshold value from the slider\n",
    "def update_plot(threshold):\n",
    "    for i, (persistence, simplex_tree, distance_matrix, tokens) in enumerate(results):\n",
    "        plot_simplicial_complex_3d(simplex_tree, distance_matrix, f\"Simplicial Complex for Image {i+1}\", threshold, tokens)\n",
    "\n",
    "# Connect the slider to the update_plot function\n",
    "widgets.interact(update_plot, threshold=threshold_slider)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29180c57-e448-4d5c-81c7-d355af6ae980",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit ('anaconda3': virtualenv)",
   "language": "python",
   "name": "python385jvsc74a57bd0474c67ce7e36ad5731492349411c4ce02ca5c170a680b2d1efe1eb0325e35fe7"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}