{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1d1f4b5b-4e05-4957-b92f-4ecee448c3cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from scipy.sparse.linalg import svds\n",
    "from tqdm.notebook import tqdm\n",
    "from random import shuffle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "263621fb-e477-42ee-a743-4b4c18ef9c79",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.backends.backend_pdf import PdfPages\n",
    "from IPython.display import Image, display, HTML\n",
    "from joblib import Parallel, delayed\n",
    "import matplotlib.image as mpimg\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "sns.set_theme()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "be8550b5-0a37-4f5e-b31f-d987c42bf9bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from numpy import dot\n",
    "from numpy.linalg import norm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d800b686-5c68-4643-9b1e-f081a91a44d3",
   "metadata": {
    "tags": []
   },
   "source": [
    "#### FUNCTIONS to load data "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "cc29b147-3e95-455d-ade5-8bec76b5866a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embeddings(model):\n",
    "    return np.load(f'../CelebA/embd_{model}_r100.npy')\n",
    "def get_names():\n",
    "    return np.load('../CelebA/names.npy')\n",
    "def get_ids(names):\n",
    "    ids_df = pd.read_csv(\"../CelebA/identity_CelebA.txt\", sep='\\s+', header=None, names=['image_id', 'id'])\n",
    "    names_ids = {os.path.basename(n):i  for i, n in enumerate(names)}\n",
    "    names_df = pd.DataFrame(names_ids.items(), columns=['image_id', 'index'])\n",
    "    names_df = names_df.merge(ids_df, on='image_id', how='left')\n",
    "    ids = names_df['id'].tolist()\n",
    "    ids = np.array(ids)\n",
    "    return ids"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "71df9b29-fc5e-4a34-878b-177af8100f51",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_direction(id_labels, embeddings):\n",
    "    # Map labels to integers and get counts in one pass.\n",
    "    unique_labels, integer_labels, counts = np.unique(id_labels, return_inverse=True, return_counts=True)\n",
    "    # Compute per-sample weights: each sample gets weight = 1 / (number of times its label occurs).\n",
    "    weights = 1 / counts[integer_labels]\n",
    "    # Compute the weighted sum using dot @\n",
    "    v = weights.astype(np.float32) @ embeddings\n",
    "    # The sum of weights is equal to the number of unique labels, so we normalize by that\n",
    "    v /= unique_labels.size\n",
    "    # Normalize the vector v\n",
    "    return v / np.linalg.norm(v)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "931f190d-b35c-4bc1-b614-17f9df7121d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def grow(seed):\n",
    "    elements = seed\n",
    "    for _ in range(1000): \n",
    "        v = get_direction(ids[elements], embeddings[elements])\n",
    "        \n",
    "        proj = embeddings @ v\n",
    "\n",
    "        for i in np.argsort(proj)[::-1]:\n",
    "            if (i not in elements):\n",
    "                if proj[i] < 0.35:\n",
    "                    return elements\n",
    "                elements = elements + [i]\n",
    "                break\n",
    "            \n",
    "    return elements"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e5223f1-9141-4d4e-b702-2d9ae6694148",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "fc638865-b2b3-4ad2-ac8a-624cb61a9059",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"partialFC\"\n",
    "# model_name = \"elasticface\"\n",
    "# model_name = \"cosface\"\n",
    "# model_name = \"arcface\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b4bfacf6-1794-4a87-81f1-5af338cc8db6",
   "metadata": {},
   "outputs": [],
   "source": [
    "embeddings = get_embeddings(model_name)\n",
    "names = get_names()\n",
    "ids = get_ids(names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "976f3889-369f-433e-b407-6ae492281964",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(202513, 512)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "embeddings.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c9978a3-8eab-4b82-9408-fcb0aef22c0e",
   "metadata": {},
   "source": [
    "### Compute scores"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "66b7fdf2-db8f-437d-b723-8d99a82d082c",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = [7735, 45242, 99111, 14216]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "1e0af932-1691-4578-86bd-b37365376aa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = [142626, 75068, 6165, 67382]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "51c69107-9e6e-46cb-b114-6ee24df61d64",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = [4712, 158511, 76055, 131633]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 541,
   "id": "c65a97c7-e068-47b4-afb3-2e04ae04a8ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = [78004, 33046, 129045, 35087, 35050, 74676]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 545,
   "id": "d609200d-5ef3-41ac-81dc-0690e3becc98",
   "metadata": {},
   "outputs": [],
   "source": [
    "seed = [15900, 165802, 119423, 34454, 141420, 12469] #176864 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "578753e7-55ab-44b9-8ca0-5fd567bad893",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div style=\"display:flex;align-items:center;justify-content:center;\"><div style=\"margin: 2px; text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/126256.jpg\" style=\"height: 200px;\"></div><div style=\"margin: 2px; text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/003905.jpg\" style=\"height: 200px;\"></div><div style=\"margin: 2px; text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/065337.jpg\" style=\"height: 200px;\"></div><div style=\"margin: 2px; text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/023660.jpg\" style=\"height: 200px;\"></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image_filenames = names[seed]\n",
    "image_filenames = np.char.add('../CelebA/', image_filenames)\n",
    "images_html = \"\".join(\n",
    "    f'<div style=\"margin: 2px; text-align: center;\"><img src=\"{image_path}\" style=\"height: 200px;\"></div>'\n",
    "    for image_path in image_filenames\n",
    ")\n",
    "html = f'<div style=\"display:flex;align-items:center;justify-content:center;\">{images_html}</div>'\n",
    "display(HTML(html))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "daf0292d-a79e-4ac6-a76d-6b271a262d79",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "554"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "elems = grow(seed)\n",
    "len(elems)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "id": "76ff7a46-665f-4276-85e8-9ecdbb170d3d",
   "metadata": {},
   "outputs": [],
   "source": [
    "pal = sns.cubehelix_palette(5, rot=-.5, light=.95)\n",
    "def rgb_to_hex(rgb):\n",
    "    return f\"#{int(rgb[0]*255):02x}{int(rgb[1]*255):02x}{int(rgb[2]*255):02x}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "e8e45c7b-2f41-41c5-b1ff-b5bf9c2b5501",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "\n",
       "<div style=\"display: flex; align-items: center; background-color: #f2f4ea; padding: 10px; border-radius: 10px;\">\n",
       "    <div style=\"display: flex; flex-wrap: wrap; justify-content: center; align-items: center; width: 100%;\">\n",
       "        <div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/003905.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/152778.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/027205.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/085164.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/113104.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/041061.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/038323.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/085507.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/028728.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/119581.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/072525.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/102836.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/055126.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/044990.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/158613.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/024109.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/127465.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/144924.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/132803.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/030675.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/136816.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/023660.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/151378.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/068900.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/092657.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/138931.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/053338.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/023036.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/128595.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/101901.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/048666.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/160206.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/079419.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/140615.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/127730.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/057642.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/015744.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/024734.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/129461.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/138424.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/067574.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/051967.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/141385.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/096943.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/042725.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/149032.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/102313.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/019525.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/138765.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/118868.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/037010.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/134839.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/127648.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/072966.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/136731.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div><div style=\"text-align: center;\"><img src=\"../ICCV 2025/CelebA/img_align_celeba/117220.jpg\" style=\"height: 150px; display: block; margin:3px;\"></div>\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "image_filenames = names[elems[:60]]\n",
    "image_filenames = np.char.add('../CelebA/', image_filenames)\n",
    "images_html = \"\".join(\n",
    "    f'<div style=\"text-align: center;\"><img src=\"{image_path}\" style=\"height: 150px; display: block; margin:3px;\"></div>'\n",
    "    for image_path in image_filenames\n",
    ")\n",
    "\n",
    "# Convert the RGB color from the palette to hex\n",
    "background_color = rgb_to_hex(pal[0])\n",
    "\n",
    "# Use flex-wrap and align items correctly\n",
    "html = f'<div style=\"display: flex; flex-wrap: wrap; justify-content: flex-start; align-items: flex-start; background-color: {background_color}; padding: 3px; border-radius: 5px;\">{images_html}</div>'\n",
    "\n",
    "# Generate the HTML for the current cluster and append it to the all_html string\n",
    "html = f'''\n",
    "<div style=\"display: flex; align-items: center; background-color: {background_color}; padding: 10px; border-radius: 10px;\">\n",
    "    <div style=\"display: flex; flex-wrap: wrap; justify-content: center; align-items: center; width: 100%;\">\n",
    "        {images_html}\n",
    "    </div>\n",
    "</div>'''\n",
    "display(HTML(html))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99be6c86-59db-4de2-be2e-44b5dbbc0998",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
