{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "entitled-exploration",
   "metadata": {},
   "source": [
    "# Image Clustering\n",
    "\n",
    "This example shows how SentenceTransformer can be used for Image Clustering.\n",
    "\n",
    "As model, we use the [OpenAI CLIP Model](https://github.com/openai/CLIP), which was trained on a large set of images and image alt texts.\n",
    "\n",
    "As a source for fotos, we use the [Unsplash Dataset Lite](https://unsplash.com/data), which contains about 25k images. See the [License](https://unsplash.com/license) about the Unsplash images. \n",
    "\n",
    "We encode all images into vector space and then find high density regions in this vector space, i.e., regions where the images are fairly similar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "acoustic-pastor",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer, util\n",
    "from PIL import Image\n",
    "import glob\n",
    "import torch\n",
    "import pickle\n",
    "import zipfile\n",
    "from IPython.display import display\n",
    "from IPython.display import Image as IPImage\n",
    "import os\n",
    "from tqdm.autonotebook import tqdm\n",
    "\n",
    "#First, we load the respective CLIP model\n",
    "model = SentenceTransformer('clip-ViT-B-32')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "major-injury",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Next, we get about 25k images from Unsplash \n",
    "img_folder = 'photos/'\n",
    "if not os.path.exists(img_folder) or len(os.listdir(img_folder)) == 0:\n",
    "    os.makedirs(img_folder, exist_ok=True)\n",
    "    \n",
    "    photo_filename = 'unsplash-25k-photos.zip'\n",
    "    if not os.path.exists(photo_filename):   #Download dataset if does not exist\n",
    "        util.http_get('http://sbert.net/datasets/'+photo_filename, photo_filename)\n",
    "        \n",
    "    #Extract all images\n",
    "    with zipfile.ZipFile(photo_filename, 'r') as zf:\n",
    "        for member in tqdm(zf.infolist(), desc='Extracting'):\n",
    "            zf.extract(member, img_folder)\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "earned-mapping",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Images: 24996\n"
     ]
    }
   ],
   "source": [
    "# Now, we need to compute the embeddings\n",
    "# To speed things up, we destribute pre-computed embeddings\n",
    "# Otherwise you can also encode the images yourself.\n",
    "# To encode an image, you can use the following code:\n",
    "# from PIL import Image\n",
    "# img_emb = model.encode(Image.open(filepath))\n",
    "\n",
    "use_precomputed_embeddings = True\n",
    "\n",
    "if use_precomputed_embeddings: \n",
    "    emb_filename = 'unsplash-25k-photos-embeddings.pkl'\n",
    "    if not os.path.exists(emb_filename):   #Download dataset if does not exist\n",
    "        util.http_get('http://sbert.net/datasets/'+emb_filename, emb_filename)\n",
    "        \n",
    "    with open(emb_filename, 'rb') as fIn:\n",
    "        img_names, img_emb = pickle.load(fIn)  \n",
    "    print(\"Images:\", len(img_names))\n",
    "else:\n",
    "    img_names = list(glob.glob('unsplash/photos/*.jpg'))\n",
    "    print(\"Images:\", len(img_names))\n",
    "    img_emb = model.encode([Image.open(filepath) for filepath in img_names], batch_size=128, convert_to_tensor=True, show_progress_bar=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "equipped-script",
   "metadata": {},
   "outputs": [],
   "source": [
    "# We have implemented our own, efficient method\n",
    "# to find high density regions in vector space\n",
    "def community_detection(embeddings, threshold, min_community_size=10, init_max_size=1000):\n",
    "    \"\"\"\n",
    "    Function for Fast Community Detection\n",
    "\n",
    "    Finds in the embeddings all communities, i.e. embeddings that are close (closer than threshold).\n",
    "\n",
    "    Returns only communities that are larger than min_community_size. The communities are returned\n",
    "    in decreasing order. The first element in each list is the central point in the community.\n",
    "    \"\"\"\n",
    "\n",
    "    # Compute cosine similarity scores\n",
    "    cos_scores = util.cos_sim(embeddings, embeddings)\n",
    "\n",
    "    # Minimum size for a community\n",
    "    top_k_values, _ = cos_scores.topk(k=min_community_size, largest=True)\n",
    "\n",
    "    # Filter for rows >= min_threshold\n",
    "    extracted_communities = []\n",
    "    for i in range(len(top_k_values)):\n",
    "        if top_k_values[i][-1] >= threshold:\n",
    "            new_cluster = []\n",
    "\n",
    "            # Only check top k most similar entries\n",
    "            top_val_large, top_idx_large = cos_scores[i].topk(k=init_max_size, largest=True)\n",
    "            top_idx_large = top_idx_large.tolist()\n",
    "            top_val_large = top_val_large.tolist()\n",
    "\n",
    "            if top_val_large[-1] < threshold:\n",
    "                for idx, val in zip(top_idx_large, top_val_large):\n",
    "                    if val < threshold:\n",
    "                        break\n",
    "\n",
    "                    new_cluster.append(idx)\n",
    "            else:\n",
    "                # Iterate over all entries (slow)\n",
    "                for idx, val in enumerate(cos_scores[i].tolist()):\n",
    "                    if val >= threshold:\n",
    "                        new_cluster.append(idx)\n",
    "\n",
    "            extracted_communities.append(new_cluster)\n",
    "\n",
    "    # Largest cluster first\n",
    "    extracted_communities = sorted(extracted_communities, key=lambda x: len(x), reverse=True)\n",
    "\n",
    "    # Step 2) Remove overlapping communities\n",
    "    unique_communities = []\n",
    "    extracted_ids = set()\n",
    "\n",
    "    for community in extracted_communities:\n",
    "        add_cluster = True\n",
    "        for idx in community:\n",
    "            if idx in extracted_ids:\n",
    "                add_cluster = False\n",
    "                break\n",
    "\n",
    "        if add_cluster:\n",
    "            unique_communities.append(community)\n",
    "            for idx in community:\n",
    "                extracted_ids.add(idx)\n",
    "\n",
    "    return unique_communities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "oriental-shower",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Total number of clusters: 147\n"
     ]
    }
   ],
   "source": [
    "# Now we run the clustering algorithm\n",
    "# With the threshold parameter, we define at which threshold we identify\n",
    "# two images as similar. Set the threshold lower, and you will get larger clusters which have \n",
    "# less similar images in it (e.g. black cat images vs. cat images vs. animal images).\n",
    "# With min_community_size, we define that we only want to have clusters of a certain minimal size\n",
    "clusters = community_detection(img_emb, threshold=0.9, min_community_size=10)\n",
    "print(\"Total number of clusters:\", len(clusters))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "buried-botswana",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 482\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 66768 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 129700 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 77884 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 401\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 35448 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 31248 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 67504 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 360\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 114944 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 64232 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 104140 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 258\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 131576 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 110000 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 114784 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 194\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 156864 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 140852 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 52808 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 145\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 30364 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 37104 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 65760 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 117\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 56072 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 63216 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 63656 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 113\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 30660 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 25120 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 127564 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 101\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 141612 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 330020 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 144492 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Cluster size: 97\n"
     ]
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 155688 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 212960 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    },
    {
     "data": {
      "image/jpeg": "[encoded data, 102696 characters]\n",
      "text/plain": [
       "<IPython.core.display.Image object>"
      ]
     },
     "metadata": {
      "image/jpeg": {
       "width": 200
      }
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Now we output the first 10 (largest) clusters\n",
    "for cluster in clusters[0:10]:\n",
    "    print(\"\\n\\nCluster size:\", len(cluster))\n",
    "    \n",
    "    #Output 3 images\n",
    "    for idx in cluster[0:3]:\n",
    "        display(IPImage(os.path.join(img_folder, img_names[idx]), width=200))\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}