{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from IPython.display import clear_output\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from owl_dreamsim_utils import eval_with_dreamsim\n",
    "\n",
    "cache_dir = None\n",
    "\n",
    "eval_dreamsim = eval_with_dreamsim(cache_dir=cache_dir, device='cuda')\n",
    "clear_output()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# functions for D&C\n",
    "def get_gt_matrix(query_dict):\n",
    "    embs = query_dict['query_emb']\n",
    "    n = len(embs)\n",
    "    gt_matrix = torch.zeros(n, n)\n",
    "    # Fill the matrix\n",
    "    for i in range(n):\n",
    "        for j in range(i, n):  # Only calculate for i <= j\n",
    "            if i == j:\n",
    "                # Diagonal: Mean of self-similarity\n",
    "                similarity = embs[i].matmul(embs[i].t())\n",
    "                gt_matrix[i, j] = similarity.mean()\n",
    "            else:\n",
    "                # Off-diagonal: Mean of inter-group similarity\n",
    "                inter_similarity = embs[i].matmul(embs[j].t())\n",
    "                mean_similarity = inter_similarity.mean()\n",
    "                gt_matrix[i, j] = mean_similarity\n",
    "                gt_matrix[j, i] = mean_similarity  # Assign to A[j, i] without recalculating\n",
    "    gt_matrix = np.array(gt_matrix)\n",
    "    return gt_matrix\n",
    "\n",
    "def sort_by_max(A):\n",
    "    sorted_rows = np.zeros_like(A)\n",
    "    used_rows = []\n",
    "\n",
    "    # Iterate over each column\n",
    "    for i in range(A.shape[1]):\n",
    "        # Find the maximum value in the i-th column that hasn't been used yet\n",
    "        max_value = -np.inf\n",
    "        max_index = -1\n",
    "        for j in range(A.shape[0]):\n",
    "            if j not in used_rows and A[j, i] > max_value:\n",
    "                max_value = A[j, i]\n",
    "                max_index = j\n",
    "\n",
    "        # Add the row with the maximum value to the sorted array\n",
    "        sorted_rows[i] = A[max_index]\n",
    "        used_rows.append(max_index)\n",
    "    return sorted_rows\n",
    "\n",
    "def gt_distance(scores, gt_matrix, ord=None):\n",
    "    if len(scores) != len(gt_matrix):\n",
    "        print(f'Count:{len(scores)}')\n",
    "        return 1.\n",
    "    # scores = [x.mean(-1).tolist() for x in np.array(scores)] # 2x2\n",
    "    # scores = np.array(scores).mean(-1)\n",
    "    tmp = []\n",
    "    for bbox_score in scores:\n",
    "        per_bbox = []\n",
    "        for ref in bbox_score:\n",
    "            per_bbox.append(np.array(ref).mean())\n",
    "        tmp.append(per_bbox)\n",
    "    \n",
    "    # scores = [np.array(x).mean(-1) for x in scores]\n",
    "    scores = np.array(tmp)\n",
    "    scores = sort_by_max(scores)\n",
    "    return np.linalg.norm(scores - gt_matrix, ord=ord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_dict = {\n",
    "    'query_name': [\"monster toy\", \"robot toy\"],\n",
    "    'query_path': [\"examples/references/monster_toy\",\n",
    "                    \"examples/references/robot_toy\"]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 22872 characters]",
      "image/png": "[encoded data, 184452 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=256x256>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GT score: 0.93\n"
     ]
    }
   ],
   "source": [
    "image = \"examples/images/good.jpg\"\n",
    "image = Image.open(image)\n",
    "scores = eval_dreamsim.score(image, query_dict, threshold=0.4, return_round=False)\n",
    "\n",
    "gt_matrix = get_gt_matrix(query_dict)\n",
    "image.resize((256, 256)).show()\n",
    "print(f\"GT score: {1 - gt_distance(scores, gt_matrix, ord=2):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 22900 characters]",
      "image/png": "[encoded data, 177156 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=256x256>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GT score: 0.86\n"
     ]
    }
   ],
   "source": [
    "image = \"examples/images/mixed1.jpg\"\n",
    "image = Image.open(image)\n",
    "scores = eval_dreamsim.score(image, query_dict, threshold=0.4, return_round=False)\n",
    "\n",
    "gt_matrix = get_gt_matrix(query_dict)\n",
    "image.resize((256, 256)).show()\n",
    "print(f\"GT score: {1 - gt_distance(scores, gt_matrix, ord=2):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 24912 characters]",
      "image/png": "[encoded data, 189916 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=256x256>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GT score: 0.31\n"
     ]
    }
   ],
   "source": [
    "image = \"examples/images/single.jpg\"\n",
    "image = Image.open(image)\n",
    "scores = eval_dreamsim.score(image, query_dict, threshold=0.4, return_round=False)\n",
    "\n",
    "gt_matrix = get_gt_matrix(query_dict)\n",
    "image.resize((256, 256)).show()\n",
    "print(f\"GT score: {1 - gt_distance(scores, gt_matrix, ord=2):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 21696 characters]",
      "image/png": "[encoded data, 171308 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=256x256>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "GT score: 0.56\n"
     ]
    }
   ],
   "source": [
    "image = \"examples/images/mixed2.jpg\"\n",
    "image = Image.open(image)\n",
    "scores = eval_dreamsim.score(image, query_dict, threshold=0.4, return_round=False)\n",
    "\n",
    "gt_matrix = get_gt_matrix(query_dict)\n",
    "image.resize((256, 256)).show()\n",
    "print(f\"GT score: {1 - gt_distance(scores, gt_matrix, ord=2):.2f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 24192 characters]",
      "image/png": "[encoded data, 188796 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=256x256>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Count:3\n",
      "GT score: 0.00\n"
     ]
    }
   ],
   "source": [
    "image = \"examples/images/count.jpg\"\n",
    "image = Image.open(image)\n",
    "scores = eval_dreamsim.score(image, query_dict, threshold=0.4, return_round=False)\n",
    "\n",
    "gt_matrix = get_gt_matrix(query_dict)\n",
    "image.resize((256, 256)).show()\n",
    "print(f\"GT score: {1 - gt_distance(scores, gt_matrix, ord=2):.2f}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ddpo-multi",
   "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.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
