{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simple image training using CLIP model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "from sentence_transformers import SentenceTransformer, losses, InputExample\n",
    "from sentence_transformers import util\n",
    "from torch.utils.data import DataLoader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load CLIP model\n",
    "model = SentenceTransformer('clip-ViT-B-32')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import zipfile\n",
    "from tqdm import tqdm\n",
    "\n",
    "# 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 1000 images\n",
    "    with zipfile.ZipFile(photo_filename, 'r') as zf:\n",
    "        for member in tqdm(zf.infolist()[:1000], desc='Extracting'):\n",
    "            zf.extract(member, img_folder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 24996 characters]",
      "image/png": "[encoded data, 271580 characters]",
      "text/plain": [
       "<PIL.JpegImagePlugin.JpegImageFile image mode=RGB size=640x433>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Load each of the images\n",
    "photos = [Image.open(os.path.join(img_folder, photo_path)) for photo_path in os.listdir(img_folder)]\n",
    "photos[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import random\n",
    "\n",
    "train_dataset = []\n",
    "for idx in range(0, len(photos), 2):\n",
    "    # We can use image pairs directly. Because our images aren't labeled, we use a random label as an example\n",
    "    train_dataset.append(InputExample(texts=[photos[idx], photos[idx + 1]], label=random.choice([0, 1])))\n",
    "    \n",
    "    # Or images and text together\n",
    "    train_dataset.append(InputExample(texts=[photos[idx], \"This is the caption\"], label=1))\n",
    "    train_dataset.append(InputExample(texts=[photos[idx], \"This is another unrelated caption\"], label=0))\n",
    "\n",
    "    # Or just texts\n",
    "    train_dataset.append(InputExample(texts=[\"This is a caption\", \"This is a similar caption\"], label=1))\n",
    "    train_dataset.append(InputExample(texts=[\"This is a caption\", \"This is an unrelated caption\"], label=0))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll create a DataLoader that batches our data and prepare a contrastive loss function\n",
    "\n",
    "train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=4)\n",
    "train_loss = losses.ContrastiveLoss(model=model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Iteration: 100%|██████████| 63/63 [00:12<00:00,  5.11it/s]\n",
      "Iteration: 100%|██████████| 63/63 [00:11<00:00,  5.72it/s]\n",
      "Iteration: 100%|██████████| 63/63 [00:11<00:00,  5.37it/s]\n",
      "Iteration: 100%|██████████| 63/63 [00:11<00:00,  5.72it/s]\n",
      "Iteration: 100%|██████████| 63/63 [00:11<00:00,  5.69it/s]\n",
      "Epoch: 100%|██████████| 5/5 [00:57<00:00, 11.43s/it]\n"
     ]
    }
   ],
   "source": [
    "# Now we can fine-tune our model on the labeled image pairs\n",
    "model.fit([(train_dataloader, train_loss)], epochs=5, show_progress_bar=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
