ModelTrainer.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d8W8DTVwEF1M"
      },
      "source": [
        "# **DONT RUN!!!** Function which fetches data for label encoding\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_2UV42z99wBz"
      },
      "outputs": [],
      "source": [
        "!pip install geopy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_TZczKla91c7"
      },
      "outputs": [],
      "source": [
        "from geopy.geocoders import Nominatim\n",
        "geolocator = Nominatim(user_agent=\"latLongToTownNameCis419\")  # Replace 'geoapiExercises' with a relevant user-agent name for your app\n",
        "\n",
        "def reverse_geocode(lat_lon_str):\n",
        "    # Perform reverse geocoding\n",
        "    location = geolocator.reverse(lat_lon_str, exactly_one=True)\n",
        "    if location:\n",
        "        address = location.raw['address']\n",
        "        city = address.get('city', '') or address.get('town', '') or address.get('village', '') or address.get('hamlet', '') or address.get('county', '')\n",
        "        county = address.get('county') or \"United Kingdom\"\n",
        "        state = address.get('state') or \"United Kingdom\"\n",
        "        country = address.get('country', '')\n",
        "        return city, county, state, country\n",
        "    else:\n",
        "        return \"No address found for this location.\"\n",
        "\n",
        "# Example coordinates\n",
        "latitude = \"57.618735\"\n",
        "longitude = \"-5.301776\"\n",
        "address = reverse_geocode(latitude + \",\" + longitude)\n",
        "print(address)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "6PldV0ok-Sk6"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')\n",
        "\n",
        "import os\n",
        "import time\n",
        "import csv\n",
        "\n",
        "# Path to the main directory\n",
        "path = '/content/drive/My Drive/ukCities/'\n",
        "\n",
        "# Get a list of all items in the main directory\n",
        "items = os.listdir(path)\n",
        "\n",
        "# Loop through each item\n",
        "for item in items:\n",
        "    item_path = os.path.join(path, item)\n",
        "    if os.path.isdir(item_path):\n",
        "        print(f\"Contents of '{item}':\")\n",
        "        # Open a CSV file to write the coordinates and corresponding city, country\n",
        "        if os.path.exists(f'{item_path}/{item}.csv'):\n",
        "          print(f\"CSV alr found for file {item}\")\n",
        "          continue\n",
        "        with open(f'{item_path}/{item}.csv', 'w', newline='') as file:\n",
        "            writer = csv.writer(file)\n",
        "            writer.writerow([\"Coords\", \"City\",\"County\",\"State\", \"Country\"])  # Writing header\n",
        "\n",
        "            coords = os.listdir(item_path)\n",
        "            for coord in coords:\n",
        "                if \"csv\" in coord:\n",
        "                  break\n",
        "                city, county, state, country = reverse_geocode(coord)\n",
        "                print(f\" - {coord}\")\n",
        "                print(city, county, state, country)\n",
        "                writer.writerow([coord, city, county, state, country])\n",
        "                time.sleep(1.25)  # Delay to simulate reverse geocoding time\n",
        "        print(f\"--> commited file for section {item}\")\n",
        "        print()  # Add a newline for better readability between directories\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Isw-GCAlXHBc"
      },
      "source": [
        "# **Start Here:** Create the Dataloader for the model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kR4kxWvT7Ciz"
      },
      "outputs": [],
      "source": [
        "from transformers import CLIPProcessor, CLIPModel"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "VV5eRfuZFPAn"
      },
      "outputs": [],
      "source": [
        "model = CLIPModel.from_pretrained(\"openai/clip-vit-base-patch32\")\n",
        "processor = CLIPProcessor.from_pretrained(\"openai/clip-vit-base-patch32\")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Wu3RDg0KrxcI"
      },
      "outputs": [],
      "source": [
        "from concurrent.futures import ThreadPoolExecutor\n",
        "from torch.utils.data import Dataset, DataLoader\n",
        "from PIL import Image, ImageFile\n",
        "import pandas as pd\n",
        "import torch\n",
        "import os\n",
        "from google.colab import drive\n",
        "from torchvision.io import read_image\n",
        "from torchvision import transforms\n",
        "from torchvision.transforms import Compose\n",
        "\n",
        "# Import the DistilBERT tokenizer\n",
        "from transformers import DistilBertTokenizer\n",
        "import time\n",
        "\n",
        "drive.mount('/content/drive')\n",
        "ImageFile.LOAD_TRUNCATED_IMAGES = True\n",
        "\n",
        "class ParallelUKClip(Dataset):\n",
        "    def __init__(self, root_dir, transform=None):\n",
        "        self.root_dir = root_dir\n",
        "        self.transform = transform\n",
        "        self.tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
        "        self.grids = []\n",
        "        self.data = self.load_data()\n",
        "\n",
        "    def load_csv_data(self, index):\n",
        "        \"\"\" Helper function to process each CSV file in a separate thread \"\"\"\n",
        "        csv_path = os.path.join(self.root_dir, str(index), f\"{index}.csv\")\n",
        "        df = pd.read_csv(csv_path)\n",
        "        data = []\n",
        "        print(\"Grid #: \", index)\n",
        "        self.grids.append(f\"Grid {index}\")\n",
        "        for _, row in df.iterrows():\n",
        "            folder_path = os.path.join(self.root_dir, str(index), row['Coords'])\n",
        "            img_files = os.listdir(folder_path)\n",
        "            img_paths = [os.path.join(folder_path, img_file) for img_file in img_files]\n",
        "\n",
        "            label = f\"Grid {index}\"\n",
        "            tokenized_label = self.tokenizer(label, return_tensors=\"pt\", padding=\"max_length\", truncation=True, max_length=512)\n",
        "            coords = row['Coords']\n",
        "            data.extend([(img_path, tokenized_label, label, coords) for img_path in img_paths])\n",
        "        return data\n",
        "\n",
        "    def load_data(self):\n",
        "        num_grids = len(os.listdir(self.root_dir))\n",
        "        with ThreadPoolExecutor() as executor:\n",
        "            results = executor.map(self.load_csv_data, range(num_grids - 1))\n",
        "        data = []\n",
        "        for result in results:\n",
        "            data.extend(result)\n",
        "        return data\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        img_path, label, grid, coords = self.data[idx]\n",
        "        image = Image.open(img_path).convert('RGB')\n",
        "        if self.transform:\n",
        "            image = self.transform(image)\n",
        "        # Tokenize the label\n",
        "        item = {}\n",
        "        item['image'] = image\n",
        "        item['label'] = label\n",
        "        # Include if running tests\n",
        "        #item['grid'] = grid\n",
        "        #item['coords'] = coords\n",
        "        #item['path'] = img_path\n",
        "        return item"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "cityPicCount = {}\n",
        "\n",
        "class ParallelUKCitiesClip(Dataset): # Dataset for UK Cities only!\n",
        "    def __init__(self, root_dir, transform=None):\n",
        "        self.topdf = pd.read_csv(\"/content/drive/My Drive/uk_cities_top_150.csv\") # changed from UK_top_50_cities.csv\n",
        "        self.hasParents = True # change to True if top 150\n",
        "        self.root_dir = root_dir\n",
        "        self.transform = transform\n",
        "        self.tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
        "        self.paths = []\n",
        "        self.cities = []\n",
        "        self.data = self.load_data()\n",
        "\n",
        "    def load_csv_data(self, index):\n",
        "        row = self.topdf.iloc[index]\n",
        "        city = \"\"\n",
        "        if not self.hasParents:\n",
        "            city = row['City']\n",
        "        else:\n",
        "          # Check if the row is empty or if 'ParentCity' is missing or not a string\n",
        "          if not pd.isna(row['ParentCity']):\n",
        "              print(\"Test\")\n",
        "              city = row['ParentCity']\n",
        "          else:\n",
        "              city = row['City']\n",
        "        self.cities.append(city)\n",
        "        totalPics = 0\n",
        "        print(\"City: \" + city +  \" \", index)\n",
        "        og_folder_path = os.path.join(self.root_dir, str(index))\n",
        "        coord_folders = os.listdir(og_folder_path)\n",
        "        coord_folder_paths = [os.path.join(og_folder_path, coord_folder) for coord_folder in coord_folders]\n",
        "\n",
        "        data = []\n",
        "        for folder_path in coord_folder_paths:\n",
        "            img_files = os.listdir(folder_path)\n",
        "            img_paths = [os.path.join(folder_path, img_file) for img_file in img_files]\n",
        "\n",
        "            label = f\"A picture of {city}\"\n",
        "            tokenized_label = self.tokenizer(label, return_tensors=\"pt\", padding=\"max_length\", truncation=True, max_length=512)\n",
        "            data.extend([(img_path, tokenized_label, city) for img_path in img_paths])\n",
        "            totalPics += len(img_paths)\n",
        "\n",
        "        if city in cityPicCount:\n",
        "          cityPicCount[city] += totalPics\n",
        "        else:\n",
        "          cityPicCount[city] = totalPics\n",
        "        return data\n",
        "\n",
        "    def load_data(self):\n",
        "        num_grids = len(os.listdir(self.root_dir))\n",
        "        with ThreadPoolExecutor() as executor:\n",
        "            results = executor.map(self.load_csv_data, range(num_grids - 1))\n",
        "        data = []\n",
        "        for result in results:\n",
        "            data.extend(result)\n",
        "        self.paths = [item[0] for item in data]\n",
        "        self.cities = list(set(self.cities)) # remove duplicates\n",
        "        return data\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.data)\n",
        "\n",
        "    def getPaths(self):\n",
        "        return self.paths\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        img_path, label, city = self.data[idx]\n",
        "        image = Image.open(img_path).convert('RGB')\n",
        "        if self.transform:\n",
        "            image = self.transform(image)\n",
        "        # Tokenize the label\n",
        "        item = {}\n",
        "        item['image'] = image\n",
        "        item['label'] = label\n",
        "        # Include if running tests\n",
        "        #item['city'] = city\n",
        "        #item['path'] = img_path\n",
        "        return item"
      ],
      "metadata": {
        "id": "YCNHaNso7o6E"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ckCDaGZ6PArL"
      },
      "source": [
        "# Encoders"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true,
        "id": "AkpthX1Gcn_G"
      },
      "outputs": [],
      "source": [
        "!pip install timm"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y28G4b8YPAbc"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torch import nn\n",
        "import torch.nn.functional as F\n",
        "import timm\n",
        "\n",
        "class ImageEncoder(nn.Module): # Encode images to a fixed size vector\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.model = timm.create_model(\n",
        "            'resnet50', pretrained=True, num_classes=0, global_pool=\"avg\"\n",
        "        )\n",
        "        for p in self.model.parameters():\n",
        "            p.requires_grad = True\n",
        "\n",
        "    def forward(self, x):\n",
        "        return self.model(x)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sltDJ5NDdi4w"
      },
      "outputs": [],
      "source": [
        "from transformers import DistilBertModel\n",
        "\n",
        "class TextEncoder(nn.Module): # Encode text to a fixed size vector\n",
        "    def __init__(self):\n",
        "        super().__init__()\n",
        "        self.model = DistilBertModel.from_pretrained(\"distilbert-base-uncased\")\n",
        "\n",
        "        for p in self.model.parameters():\n",
        "            p.requires_grad = True\n",
        "\n",
        "        # we are using the CLS token hidden representation as the sentence's embedding\n",
        "        self.target_token_idx = 0\n",
        "\n",
        "    def forward(self, input_ids, attention_mask):\n",
        "        input_ids = input_ids.squeeze(1)\n",
        "        attention_mask = attention_mask.squeeze(1) # need to squeeze from [batch_size, 1, 512] to [batch_size, 512]\n",
        "\n",
        "        output = self.model(input_ids=input_ids, attention_mask=attention_mask)\n",
        "        last_hidden_state = output.last_hidden_state\n",
        "        return last_hidden_state[:, self.target_token_idx, :]"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Dgom9PMfgA7q"
      },
      "outputs": [],
      "source": [
        "class ProjectionHead(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        embedding_dim,\n",
        "        projection_dim=256, # required by CLIP paper\n",
        "        dropout=0.1\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.projection = nn.Linear(embedding_dim, projection_dim)\n",
        "        self.gelu = nn.GELU()\n",
        "        self.fc = nn.Linear(projection_dim, projection_dim)\n",
        "        self.dropout = nn.Dropout(dropout)\n",
        "        self.layer_norm = nn.LayerNorm(projection_dim)\n",
        "\n",
        "    def forward(self, x):\n",
        "        projected = self.projection(x)\n",
        "        x = self.gelu(projected)\n",
        "        x = self.fc(x)\n",
        "        x = self.dropout(x)\n",
        "        x = x + projected\n",
        "        x = self.layer_norm(x)\n",
        "        return x\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XeMdCCKsGazZ"
      },
      "source": [
        "# Model Architecture"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vzpnr7OLq0zr"
      },
      "outputs": [],
      "source": [
        "class CLIPModel(nn.Module):\n",
        "    def __init__(\n",
        "        self,\n",
        "        temperature=1.0,\n",
        "        image_embedding=2048,\n",
        "        text_embedding=768,\n",
        "    ):\n",
        "        super().__init__()\n",
        "        self.image_encoder = ImageEncoder()\n",
        "        self.text_encoder = TextEncoder()\n",
        "        self.image_projection = ProjectionHead(embedding_dim=image_embedding)\n",
        "        self.text_projection = ProjectionHead(embedding_dim=text_embedding)\n",
        "        self.temperature = temperature\n",
        "\n",
        "    def forward(self, batch):\n",
        "        # Getting Image and Text Features\n",
        "        image_features = self.image_encoder(batch[\"image\"])\n",
        "        text_features = self.text_encoder(\n",
        "            input_ids=batch['label'][\"input_ids\"], attention_mask=batch['label'][\"attention_mask\"]\n",
        "        )\n",
        "        # Getting Image and Text Embeddings (with same dimension)\n",
        "        image_embeddings = self.image_projection(image_features)\n",
        "        text_embeddings = self.text_projection(text_features)\n",
        "\n",
        "        # Calculating the Loss\n",
        "        logits = (text_embeddings @ image_embeddings.T) / self.temperature\n",
        "        images_similarity = image_embeddings @ image_embeddings.T\n",
        "        texts_similarity = text_embeddings @ text_embeddings.T\n",
        "        targets = F.softmax(\n",
        "            (images_similarity + texts_similarity) / 2 * self.temperature, dim=-1\n",
        "        )\n",
        "\n",
        "        texts_loss = cross_entropy(logits, targets, reduction='none')\n",
        "        images_loss = cross_entropy(logits.T, targets.T, reduction='none')\n",
        "        loss =  (images_loss + texts_loss) / 2.0\n",
        "        return loss.mean()\n",
        "\n",
        "\n",
        "def cross_entropy(preds, targets, reduction='none'):\n",
        "    log_softmax = nn.LogSoftmax(dim=-1)\n",
        "    loss = (-targets * log_softmax(preds)).sum(1)\n",
        "    if reduction == \"none\":\n",
        "        return loss\n",
        "    elif reduction == \"mean\":\n",
        "        return loss.mean()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "STU1W7QwuXba"
      },
      "source": [
        "# Train Helpers"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UKrM2918uXIs"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import albumentations as A\n",
        "from torch.utils.data import random_split\n",
        "\n",
        "def get_transforms():\n",
        "    return Compose([\n",
        "          transforms.ToTensor(),\n",
        "          transforms.Normalize([0, 0, 0], [1, 1, 1])\n",
        "        ])\n",
        "\n",
        "def build_loaders():\n",
        "    transformer = get_transforms()\n",
        "    dataset = ParallelUKCitiesClip(root_dir='/content/drive/My Drive/ukCitiesTop50/', transform=transformer) # changed from ukCitiesTop50\n",
        "\n",
        "    # Define the split sizes\n",
        "    train_split = 0.7\n",
        "    val_split = 0.2\n",
        "    test_split = 0.1\n",
        "\n",
        "    train_size = int(len(dataset) * train_split)\n",
        "    val_size = int(len(dataset) * val_split)\n",
        "    test_size = len(dataset) - train_size - val_size\n",
        "\n",
        "    # Randomly split the dataset\n",
        "    train_dataset, val_dataset, test_dataset = random_split(dataset, [train_size, val_size, test_size])\n",
        "\n",
        "    # Create data loaders for Train/Val/Test\n",
        "    train_loader = DataLoader(train_dataset, batch_size=32, num_workers=10, pin_memory=True, shuffle=True)\n",
        "    val_loader = DataLoader(val_dataset, batch_size=32, num_workers=10, pin_memory=True, shuffle=False)\n",
        "    test_loader = DataLoader(test_dataset, batch_size=32, num_workers=10, pin_memory=True, shuffle=False)\n",
        "\n",
        "    return train_loader, val_loader, test_loader, dataset\n",
        "\n",
        "class AvgMeter:\n",
        "    def __init__(self, name=\"Metric\"):\n",
        "        self.name = name\n",
        "        self.reset()\n",
        "\n",
        "    def reset(self):\n",
        "        self.avg, self.sum, self.count = [0] * 3\n",
        "\n",
        "    def update(self, val, count=1):\n",
        "        self.count += count\n",
        "        self.sum += val * count\n",
        "        self.avg = self.sum / self.count\n",
        "\n",
        "    def __repr__(self):\n",
        "        text = f\"{self.name}: {self.avg:.4f}\"\n",
        "        return text\n",
        "\n",
        "def get_lr(optimizer):\n",
        "    for param_group in optimizer.param_groups:\n",
        "        return param_group[\"lr\"]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AFtieNQm3IL5"
      },
      "source": [
        "# Train the model!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "m3JWfU8RPmKL",
        "collapsed": true
      },
      "outputs": [],
      "source": [
        "train_loader, valid_loader, test_loader, dataset = build_loaders() # Create DataLoaders. Also get Dataset (needed later)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LlTMMwn03Hh8"
      },
      "outputs": [],
      "source": [
        "import itertools\n",
        "from tqdm.autonotebook import tqdm\n",
        "from google.colab import files\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "\n",
        "def train_epoch(model, train_loader, optimizer, lr_scheduler, step):\n",
        "    loss_meter = AvgMeter()\n",
        "    tqdm_object = tqdm(train_loader, total=len(train_loader))\n",
        "    for batch in tqdm_object:\n",
        "        batch = {k: v.to(device) for k, v in batch.items() if k != \"caption\"}\n",
        "        loss = model(batch)\n",
        "        optimizer.zero_grad()\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        if step == \"batch\":\n",
        "            lr_scheduler.step()\n",
        "\n",
        "        count = batch[\"image\"].size(0)\n",
        "        loss_meter.update(loss.item(), count)\n",
        "\n",
        "        tqdm_object.set_postfix(train_loss=loss_meter.avg, lr=get_lr(optimizer))\n",
        "    return loss_meter\n",
        "\n",
        "\n",
        "def valid_epoch(model, valid_loader):\n",
        "    loss_meter = AvgMeter()\n",
        "\n",
        "    tqdm_object = tqdm(valid_loader, total=len(valid_loader))\n",
        "    for batch in tqdm_object:\n",
        "        batch = {k: v.to(device) for k, v in batch.items() if k != \"caption\"}\n",
        "        loss = model(batch)\n",
        "\n",
        "        count = batch[\"image\"].size(0)\n",
        "        loss_meter.update(loss.item(), count)\n",
        "\n",
        "        tqdm_object.set_postfix(valid_loss=loss_meter.avg)\n",
        "    return loss_meter\n",
        "\n",
        "\n",
        "def main():\n",
        "    model = CLIPModel().to(device)\n",
        "\n",
        "    # Some hyperparameters are specified in CLIP paper. We found 6 Epochs to work best.\n",
        "    image_encoder_lr = 1e-4\n",
        "    text_encoder_lr = 1e-5\n",
        "    head_lr = 1e-3\n",
        "    weight_decay = 1e-3\n",
        "    patience = 1\n",
        "    factor = 0.8\n",
        "    epochs = 6\n",
        "\n",
        "    params = [\n",
        "        {\"params\": model.image_encoder.parameters(), \"lr\": image_encoder_lr},\n",
        "        {\"params\": model.text_encoder.parameters(), \"lr\": text_encoder_lr},\n",
        "        {\"params\": itertools.chain(\n",
        "            model.image_projection.parameters(), model.text_projection.parameters()\n",
        "        ), \"lr\": head_lr, \"weight_decay\": weight_decay}\n",
        "    ]\n",
        "    optimizer = torch.optim.AdamW(params, weight_decay=0.)\n",
        "    lr_scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(\n",
        "        optimizer, mode=\"min\", patience=patience, factor=factor\n",
        "    )\n",
        "    step = \"epoch\"\n",
        "\n",
        "    best_loss = float('inf')\n",
        "    train_losses = []\n",
        "    valid_losses = []\n",
        "\n",
        "    for epoch in range(epochs):\n",
        "        print(f\"Epoch: {epoch + 1}\")\n",
        "        model.train()\n",
        "        train_loss = train_epoch(model, train_loader, optimizer, lr_scheduler, step)\n",
        "        model.eval()\n",
        "        with torch.no_grad():\n",
        "            valid_loss = valid_epoch(model, valid_loader)\n",
        "\n",
        "        train_losses.append(train_loss.avg)\n",
        "        valid_losses.append(valid_loss.avg)\n",
        "\n",
        "        if valid_loss.avg < best_loss:\n",
        "            best_loss = valid_loss.avg\n",
        "            torch.save(model.state_dict(), \"best_uk150.pt\")\n",
        "            print(\"Saved Best Model!\")\n",
        "\n",
        "        lr_scheduler.step(valid_loss.avg)\n",
        "\n",
        "    # Create plot of train vs. val loss:\n",
        "    plt.figure(figsize=(10, 5))\n",
        "    plt.plot(range(1, epochs + 1), train_losses, label='Training Loss')\n",
        "    plt.plot(range(1, epochs + 1), valid_losses, label='Validation Loss')\n",
        "    plt.xlabel('Epochs')\n",
        "    plt.ylabel('Loss')\n",
        "    plt.title('Training and Validation Loss Over Epochs')\n",
        "    plt.legend()\n",
        "    plt.show()"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hE8OQJf4VFjY"
      },
      "outputs": [],
      "source": [
        "# Distribution of images:\n",
        "print(cityPicCount)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "EmrfloOEPd-4"
      },
      "outputs": [],
      "source": [
        "# Run the model:\n",
        "main()"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Accuracy of Test Set - UK Cities 50&150 Datasets"
      ],
      "metadata": {
        "id": "OdarQS-bENqg"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "from geopy.distance import geodesic\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "from transformers import DistilBertTokenizer\n",
        "\n",
        "def generate_predictions(model, test_loader, maxBatches):\n",
        "    model.eval()\n",
        "    all_predictions = []\n",
        "    all_true_labels = []\n",
        "    all_image_paths = []\n",
        "\n",
        "    with torch.no_grad():\n",
        "        i = 0\n",
        "        for batch in test_loader:\n",
        "            if i > maxBatches:\n",
        "                break\n",
        "            images = batch['image'].to(device)\n",
        "            true_labels = batch['city']\n",
        "            image_paths = batch['path']  # Assuming 'image_path' contains the paths\n",
        "\n",
        "            # Forward pass\n",
        "            image_features = model.image_encoder(images)\n",
        "            image_embeddings = model.image_projection(image_features)\n",
        "\n",
        "            # Normalize the embeddings\n",
        "            image_embeddings = F.normalize(image_embeddings, p=2, dim=-1)\n",
        "\n",
        "            # Store the results\n",
        "            all_predictions.extend(image_embeddings.cpu().numpy())\n",
        "            all_true_labels.extend(true_labels)\n",
        "            all_image_paths.extend(image_paths)\n",
        "            i+= 1\n",
        "\n",
        "    return np.array(all_predictions), all_true_labels, all_image_paths\n",
        "\n",
        "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "model = CLIPModel().to(device)\n",
        "model.load_state_dict(torch.load(\"drive/MyDrive/best_uk150.pt\", map_location=device))\n",
        "\n",
        "predictions, true_labels, image_paths = generate_predictions(model, test_loader, 50)\n",
        "\n",
        "def calculate_accuracy(model, predictions, true_labels, cities):\n",
        "    cities_embeddings = []\n",
        "\n",
        "    # Generate embeddings for all city labels\n",
        "    for city in cities:\n",
        "        label = f\"A picture of {city}\"\n",
        "        tokenized_label = tokenizer(label, return_tensors=\"pt\", padding=\"max_length\", truncation=True, max_length=512)\n",
        "        with torch.no_grad():\n",
        "            text_features = model.text_encoder(\n",
        "                input_ids=tokenized_label[\"input_ids\"].to(device),\n",
        "                attention_mask=tokenized_label[\"attention_mask\"].to(device)\n",
        "            )\n",
        "            text_embedding = model.text_projection(text_features)\n",
        "            text_embedding = F.normalize(text_embedding, p=2, dim=-1)\n",
        "            cities_embeddings.append(text_embedding.cpu().numpy())\n",
        "\n",
        "    cities_embeddings = np.vstack(cities_embeddings)\n",
        "    city_indices = {city: idx for idx, city in enumerate(cities)}\n",
        "\n",
        "    correct_predictions = 0\n",
        "    total_predictions = len(predictions)\n",
        "    total_distance = 0\n",
        "    distances = []\n",
        "    topdf = pd.read_csv(\"/content/drive/My Drive/uk_cities_top_150.csv\") # changed from top 50\n",
        "    predicted_labels = []\n",
        "\n",
        "    for i, img_embedding in enumerate(predictions):\n",
        "        similarities = np.dot(cities_embeddings, img_embedding)\n",
        "        predicted_index = np.argmax(similarities)\n",
        "        predicted_city = cities[predicted_index]\n",
        "        true_city = true_labels[i]\n",
        "\n",
        "        predicted_labels.append(predicted_city)\n",
        "\n",
        "        if predicted_city == true_city:\n",
        "            correct_predictions += 1\n",
        "        dist = distance(predicted_city, true_city, topdf)\n",
        "        total_distance += dist\n",
        "        distances.append(dist)\n",
        "\n",
        "    accuracy = correct_predictions / total_predictions\n",
        "    average_distance = total_distance / total_predictions\n",
        "    thresholds = [25, 75, 150]\n",
        "    out_percentages = []\n",
        "    for threshold in thresholds:\n",
        "        total_in_threshold = 0\n",
        "        for dist in distances:\n",
        "            if dist < threshold:\n",
        "                total_in_threshold+=1\n",
        "        out_percentages.append((threshold, total_in_threshold/len(distances)))\n",
        "\n",
        "    return accuracy, average_distance, predicted_labels, out_percentages\n",
        "\n",
        "def get_coordinates(town, df):\n",
        "    row = df[df[\"City\"] == town]\n",
        "    if not row.empty:\n",
        "        return (row.iloc[0][\"Latitude\"], row.iloc[0][\"Longitude\"])\n",
        "    else:\n",
        "        raise ValueError(f\"Town '{town}' not found in the DataFrame.\")\n",
        "\n",
        "def distance(town1, town2, df):\n",
        "    coord1 = get_coordinates(town1, df)\n",
        "    coord2 = get_coordinates(town2, df)\n",
        "    return geodesic(coord1, coord2).kilometers\n",
        "\n",
        "cities = dataset.cities\n",
        "accuracy, average_distance, predicted_labels, percentages = calculate_accuracy(model, predictions, true_labels, cities)\n",
        "print(f\"Test Dataset (10%) Accuracy: {accuracy * 100:.2f}%\")\n",
        "print(f\"Average Distance: {average_distance} KM\")\n",
        "print(f'Number of Predictions: {len(predictions)}')\n",
        "#print(f'Number of Unique Cities After Using Parent City: {len(cities)}')\n",
        "for (threshold, percentage) in percentages:\n",
        "    print(f\"Percentage of predictions within {threshold} KM: {percentage * 100:.2f}%\")\n",
        "\n",
        "# Select 5 random images\n",
        "indices = random.sample(range(len(image_paths)), 5)\n",
        "\n",
        "# Plot the images with their predicted and true labels\n",
        "fig, axes = plt.subplots(1, 5, figsize=(20, 10))\n",
        "for idx, ax in zip(indices, axes):\n",
        "    image_path = image_paths[idx]\n",
        "    image = plt.imread(image_path)\n",
        "    predicted_label = predicted_labels[idx]\n",
        "    true_label = true_labels[idx]\n",
        "\n",
        "    ax.imshow(image)\n",
        "    ax.set_title(f\"Predicted: {predicted_label}\\nTrue: {true_label}\")\n",
        "    ax.axis('off')\n",
        "\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "DOkuSJ21EQsH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "# Accuracy of UK Grid Model"
      ],
      "metadata": {
        "id": "Z2ZlJZMapKAN"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install shapely==1.8"
      ],
      "metadata": {
        "id": "qkRg2mRF3ma6"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch.nn.functional as F\n",
        "import numpy as np\n",
        "from geopy.distance import geodesic\n",
        "import pandas as pd\n",
        "import matplotlib.pyplot as plt\n",
        "import random\n",
        "from transformers import DistilBertTokenizer\n",
        "import pickle\n",
        "from shapely.geometry import Polygon\n",
        "\n",
        "# Get center of each grid in coordinates:\n",
        "\n",
        "Ukpolygrid = pickle.load(open(\"/content/drive/My Drive/ukPolyGrid.pkl\",'rb'))\n",
        "grid_center_coords = []\n",
        "for grid,coor in Ukpolygrid.items():\n",
        "    poly = Polygon(np.flip(coor))\n",
        "    minx, miny, maxx, maxy = poly.bounds\n",
        "    centerx = (minx + maxx) / 2\n",
        "    centery = (miny + maxy) / 2\n",
        "    grid_center_coords.append((centery, centerx))\n",
        "\n",
        "def generate_predictions(model, test_loader, maxBatches):\n",
        "    model.eval()\n",
        "    all_predictions = []\n",
        "    all_true_labels = []\n",
        "    all_image_paths = []\n",
        "    all_image_coords = []\n",
        "\n",
        "    with torch.no_grad():\n",
        "        i = 0\n",
        "        for batch in test_loader:\n",
        "            if i > maxBatches:\n",
        "                break\n",
        "            images = batch['image'].to(device)\n",
        "            true_labels = batch['grid']\n",
        "            image_paths = batch['path']  # Assuming 'image_path' contains the paths\n",
        "            coords = batch['coords']\n",
        "\n",
        "            # Forward pass\n",
        "            image_features = model.image_encoder(images)\n",
        "            image_embeddings = model.image_projection(image_features)\n",
        "\n",
        "            # Normalize the embeddings\n",
        "            image_embeddings = F.normalize(image_embeddings, p=2, dim=-1)\n",
        "\n",
        "            # Store the results\n",
        "            all_predictions.extend(image_embeddings.cpu().numpy())\n",
        "            all_true_labels.extend(true_labels)\n",
        "            all_image_paths.extend(image_paths)\n",
        "            all_image_coords.extend(coords)\n",
        "            i+= 1\n",
        "\n",
        "    return np.array(all_predictions), all_true_labels, all_image_paths, all_image_coords\n",
        "\n",
        "tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')\n",
        "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
        "model = CLIPModel().to(device)\n",
        "model.load_state_dict(torch.load(\"drive/MyDrive/best_uk_grid.pt\", map_location=device))\n",
        "\n",
        "predictions, true_labels, image_paths, image_coords = generate_predictions(model, test_loader, 50)\n",
        "\n",
        "def calculate_accuracy(model, predictions, true_labels, grids, image_coords):\n",
        "    grid_embeddings = []\n",
        "\n",
        "    # Generate embeddings for all grid labels\n",
        "    for grid in grids:\n",
        "        label = grid\n",
        "        tokenized_label = tokenizer(label, return_tensors=\"pt\", padding=\"max_length\", truncation=True, max_length=512)\n",
        "        with torch.no_grad():\n",
        "            text_features = model.text_encoder(\n",
        "                input_ids=tokenized_label[\"input_ids\"].to(device),\n",
        "                attention_mask=tokenized_label[\"attention_mask\"].to(device)\n",
        "            )\n",
        "            text_embedding = model.text_projection(text_features)\n",
        "            text_embedding = F.normalize(text_embedding, p=2, dim=-1)\n",
        "            grid_embeddings.append(text_embedding.cpu().numpy())\n",
        "\n",
        "    grid_embeddings = np.vstack(grid_embeddings)\n",
        "    grid_indices = {grid: idx for idx, grid in enumerate(grid)}\n",
        "\n",
        "    correct_predictions = 0\n",
        "    total_predictions = len(predictions)\n",
        "    total_distance = 0\n",
        "    distances = []\n",
        "    topdf = pd.read_csv(\"/content/drive/My Drive/uk_cities_top_150.csv\") # changed from top 50\n",
        "    predicted_labels = []\n",
        "\n",
        "    for i, img_embedding in enumerate(predictions):\n",
        "        similarities = np.dot(grid_embeddings, img_embedding)\n",
        "        predicted_index = np.argmax(similarities)\n",
        "\n",
        "        predicted_grid = grids[predicted_index]\n",
        "        true_grid = true_labels[i]\n",
        "\n",
        "        predicted_coords = image_coords[predicted_index]\n",
        "        true_grid_number = int(true_grid[4:])\n",
        "        true_coords = grid_center_coords[true_grid_number] # Given grid index i, we get the center of that grid\n",
        "\n",
        "        predicted_labels.append(predicted_grid)\n",
        "\n",
        "        if predicted_grid == true_grid:\n",
        "            correct_predictions += 1\n",
        "\n",
        "        dist = distance(predicted_coords, true_coords)\n",
        "        total_distance += dist\n",
        "        distances.append(dist)\n",
        "\n",
        "    accuracy = correct_predictions / total_predictions\n",
        "    average_distance = total_distance / total_predictions\n",
        "    thresholds = [25, 75, 150]\n",
        "    out_percentages = []\n",
        "    for threshold in thresholds:\n",
        "        total_in_threshold = 0\n",
        "        for dist in distances:\n",
        "            if dist < threshold:\n",
        "                total_in_threshold+=1\n",
        "        out_percentages.append((threshold, total_in_threshold/len(distances)))\n",
        "\n",
        "    return accuracy, average_distance, predicted_labels, out_percentages\n",
        "\n",
        "def distance(coord1, coord2):\n",
        "    return geodesic(coord1, coord2).kilometers\n",
        "\n",
        "grids = dataset.grids\n",
        "accuracy, average_distance, predicted_labels, percentages = calculate_accuracy(model, predictions, true_labels, grids, image_coords)\n",
        "print(f\"Test Dataset (10%) Accuracy: {accuracy * 100:.2f}%\")\n",
        "print(f\"Average Distance: {average_distance} KM\")\n",
        "print(f'Number of Predictions: {len(predictions)}')\n",
        "#print(f'Number of Unique Cities After Using Parent City: {len(cities)}')\n",
        "for (threshold, percentage) in percentages:\n",
        "    print(f\"Percentage of predictions within {threshold} KM: {percentage * 100:.2f}%\")\n",
        "\n",
        "# Select 5 random images\n",
        "indices = random.sample(range(len(image_paths)), 5)\n",
        "\n",
        "# Plot the images with their predicted and true labels\n",
        "fig, axes = plt.subplots(1, 5, figsize=(20, 10))\n",
        "for idx, ax in zip(indices, axes):\n",
        "    image_path = image_paths[idx]\n",
        "    image = plt.imread(image_path)\n",
        "    predicted_label = predicted_labels[idx]\n",
        "    true_label = true_labels[idx]\n",
        "\n",
        "    ax.imshow(image)\n",
        "    ax.set_title(f\"Predicted: {predicted_label}\\nTrue: {true_label}\")\n",
        "    ax.axis('off')\n",
        "\n",
        "plt.show()\n"
      ],
      "metadata": {
        "id": "gDBLUvKlpPHH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6Ak5ggiPlxcQ"
      },
      "source": [
        "# Code to clear GPU RAM\n",
        "\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4yzCqBLai3cd"
      },
      "outputs": [],
      "source": [
        "import gc\n",
        "\n",
        "model = None\n",
        "gc.collect()\n",
        "torch.cuda.empty_cache()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AV9uw6KkSKzC"
      },
      "source": [
        "# Sources:\n",
        "\n",
        "https://arxiv.org/pdf/2302.00275\n",
        "\n",
        "https://github.com/moein-shariatnia/OpenAI-CLIP"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [
        "Z2ZlJZMapKAN"
      ],
      "gpuType": "L4",
      "machine_shape": "hm",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}