CirclesYOLOv8.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import transforms\n",
    "import torchvision.transforms as T\n",
    "import torch.nn.functional as F\n",
    "from torchvision.io import read_image, ImageReadMode\n",
    "from PIL import Image, ImageDraw, ImageFilter\n",
    "import os\n",
    "import json\n",
    "import pickle\n",
    "import numpy as np\n",
    "\n",
    "class RefCocoG_Dataset(Dataset):\n",
    "    full_annotations = None\n",
    "\n",
    "    def __init__(self, root_dir, annotations_f, instances_f, split='train', transform=None, target_transform=None) -> None:\n",
    "        super().__init__()\n",
    "\n",
    "        self.root_dir = root_dir\n",
    "        self.annotations_f = annotations_f\n",
    "        self.instances_f = instances_f\n",
    "\n",
    "        self.split = split\n",
    "\n",
    "        self.transform = transform\n",
    "        self.target_transform = target_transform\n",
    "\n",
    "        self.get_annotations()\n",
    "        self.image_names = list([\n",
    "            self.annotations[id]['image']['actual_file_name']\n",
    "            for id in self.annotations\n",
    "        ])\n",
    "\n",
    "    def get_annotations(self):\n",
    "        if RefCocoG_Dataset.full_annotations:\n",
    "            self.annotations = dict(filter(lambda match: match[1]['image']['split'] == self.split, RefCocoG_Dataset.full_annotations.items()))\n",
    "            return\n",
    "\n",
    "        # Load pickle data\n",
    "        with open(os.path.join(self.root_dir, 'annotations', self.annotations_f), 'rb') as file:\n",
    "            self.data = pickle.load(file)\n",
    "\n",
    "        # Load instances\n",
    "        with open(os.path.join(self.root_dir, 'annotations', self.instances_f), 'rb') as file:\n",
    "            self.instances = json.load(file)\n",
    "\n",
    "        # Match data between the two files and build the actual dataset\n",
    "        self.annotations = {}\n",
    "\n",
    "        images_actual_file_names = {}\n",
    "        for image in self.instances['images']:\n",
    "            images_actual_file_names[image['id']] = image['file_name']\n",
    "\n",
    "        for image in self.data:\n",
    "            if image['ann_id'] not in self.annotations:\n",
    "                self.annotations[image['ann_id']] = {}\n",
    "\n",
    "            self.annotations[image['ann_id']]['image'] = image\n",
    "            self.annotations[image['ann_id']]['image']['actual_file_name'] = images_actual_file_names[image['image_id']]\n",
    "\n",
    "        for annotation in self.instances['annotations']:\n",
    "            if annotation['id'] not in self.annotations:\n",
    "                continue\n",
    "\n",
    "            self.annotations[annotation['id']]['annotation'] = annotation\n",
    "\n",
    "        # Keep only samples from the given split\n",
    "        RefCocoG_Dataset.full_annotations = self.annotations\n",
    "        self.annotations = dict(filter(lambda match: match[1]['image']['split'] == self.split, self.annotations.items()))\n",
    "\n",
    "    def __len__(self):\n",
    "        # Return the number of images\n",
    "        return len(self.image_names)\n",
    "\n",
    "    def corner_size_to_corners(self, bounding_box):\n",
    "        \"\"\"\n",
    "        Transform (top_left_x, top_left_y, width, height) bounding box representation\n",
    "        into (top_left_x, top_left_y, bottom_right_x, bottom_right_y)\n",
    "        \"\"\"\n",
    "\n",
    "        return [\n",
    "            bounding_box[0],\n",
    "            bounding_box[1],\n",
    "            bounding_box[0] + bounding_box[2],\n",
    "            bounding_box[1] + bounding_box[3]\n",
    "        ]\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        # Get the image name at the given index\n",
    "        image_name = self.image_names[idx]\n",
    "\n",
    "        # Load the image file as a PIL image\n",
    "        image = Image.open(os.path.join(self.root_dir, 'images', image_name)).convert('RGB')\n",
    "        # image = read_image(os.path.join(self.root_dir, 'images', image_name), ImageReadMode.RGB)\n",
    "        \n",
    "        image_id = list(self.annotations)[idx]\n",
    "\n",
    "        # print(image_id)\n",
    "\n",
    "        # Get the caption for the image\n",
    "        prompts = [\n",
    "            prompt['sent'] for prompt in self.annotations[image_id]['image']['sentences']\n",
    "        ]\n",
    "\n",
    "        # Get the bounding box for the prompts for the image\n",
    "        bounding_box = self.corner_size_to_corners(self.annotations[image_id]['annotation']['bbox'])\n",
    "\n",
    "        # Apply the transform if given\n",
    "        if self.transform:\n",
    "            image = self.transform(image)\n",
    "\n",
    "        sample = [\n",
    "            image,\n",
    "            bounding_box,\n",
    "            prompts,\n",
    "        ]\n",
    "\n",
    "        # Return the sample as a list\n",
    "        return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_train = RefCocoG_Dataset('refcocog', 'refs(umd).p', 'instances.json', split='train')\n",
    "dataset_val = RefCocoG_Dataset('refcocog', 'refs(umd).p', 'instances.json', split='val')\n",
    "dataset_test = RefCocoG_Dataset('refcocog', 'refs(umd).p', 'instances.json', split='test')\n",
    "\n",
    "dataset_splits = [\n",
    "    dataset_train,\n",
    "    dataset_val,\n",
    "    dataset_test\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(RefCocoG_Dataset.full_annotations), len(dataset_train.annotations), len(dataset_val.annotations), len(dataset_test.annotations)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def collate_differently_sized_prompts(batch):\n",
    "    images = [item[0] for item in batch]\n",
    "    bboxes = [item[1] for item in batch]\n",
    "    prompts = [item[2] for item in batch]\n",
    "    \n",
    "    return list(images), list(bboxes), list(prompts)\n",
    "\n",
    "def get_data(dataset_splits, batch_size=64, test_batch_size=256, num_workers=0):\n",
    "    training_data = dataset_splits[0]\n",
    "    validation_data = dataset_splits[1]\n",
    "    test_data = dataset_splits[2]\n",
    "\n",
    "    # Change shuffle to True for train\n",
    "    train_loader = torch.utils.data.DataLoader(training_data, batch_size, shuffle=True, drop_last=True, collate_fn=collate_differently_sized_prompts, num_workers=num_workers)\n",
    "    val_loader = torch.utils.data.DataLoader(validation_data, test_batch_size, shuffle=False, collate_fn=collate_differently_sized_prompts, num_workers=num_workers)\n",
    "    test_loader = torch.utils.data.DataLoader(test_data, test_batch_size, shuffle=False, collate_fn=collate_differently_sized_prompts, num_workers=num_workers)\n",
    "\n",
    "    return train_loader, val_loader, test_loader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader, val_loader, test_loader = get_data(dataset_splits, batch_size=64, test_batch_size=64, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:0\") # First GPU\n",
    "else:\n",
    "    device = 'cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from ultralytics import YOLO\n",
    "\n",
    "if torch.cuda.is_available():\n",
    "    yolo_models = [YOLO('yolov8s.pt') for i in range(torch.cuda.device_count())]\n",
    "else:\n",
    "    yolo_models = [YOLO('yolov8s.pt')]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import clip\n",
    "\n",
    "# clip_backbone = 'ViT-B/32'\n",
    "# clip_backbone = 'ViT-L/14'\n",
    "# clip_backbone = 'ViT-L/14@336px'\n",
    "# clip_backbone = 'RN50x16'\n",
    "# clip_backbone = 'RN50x64'\n",
    "# clip_backbone = 'RN101'\n",
    "\n",
    "clip_backbones = [\n",
    "    # 'RN50x64'\n",
    "    # 'RN50x16',\n",
    "    # 'ViT-B/16'\n",
    "    'ViT-L/14@336px',\n",
    "    # 'ViT-B/32'\n",
    "]\n",
    "\n",
    "models, preprocesses = {}, {}\n",
    "\n",
    "for clip_backbone in clip_backbones:\n",
    "    models[clip_backbone] = []\n",
    "    preprocesses[clip_backbone] = []\n",
    "\n",
    "    if torch.cuda.is_available():\n",
    "        for i in range(torch.cuda.device_count()):\n",
    "            model, preprocess = clip.load(clip_backbone, device=f'cuda:{i}')\n",
    "            \n",
    "            models[clip_backbone].append(model)\n",
    "            preprocesses[clip_backbone].append(preprocess)\n",
    "    else:\n",
    "        model, preprocess = clip.load(clip_backbone, device=device)\n",
    "        models[clip_backbone].append(model)\n",
    "        preprocesses[clip_backbone].append(preprocess)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# next(models['RN50x16'][0].parameters()).device"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def cosine_similarity(a: torch.Tensor, b: torch.Tensor):\n",
    "    \"\"\"\n",
    "    Cosine Similarity\n",
    "\n",
    "    Normalizes both tensors a and b. Returns <b, a.T> (inner product).\n",
    "    \"\"\"\n",
    "\n",
    "    a_norm = a / a.norm(dim=-1, keepdim=True)\n",
    "    b_norm = b / b.norm(dim=-1, keepdim=True)\n",
    "\n",
    "    similarity = (b_norm @ a_norm.T)\n",
    "\n",
    "    return similarity.cpu()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "def visualise_scores(scores: torch.Tensor, images, texts: list[str]):\n",
    "    for t_idx, text in enumerate(texts):\n",
    "        for i_idx, image in enumerate(images):\n",
    "            fig, ax = plt.subplots()\n",
    "            ax.imshow(image)\n",
    "            ax.set_title(f'Score: {scores[t_idx, i_idx]} / Prompt: {text}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_prompts = []\n",
    "\n",
    "for batch_idx, (images, gt_bounding_boxes, prompts) in enumerate(train_loader):\n",
    "    refined_prompts = ['This is ' + prompt for sample in prompts for prompt in sample]\n",
    "\n",
    "    for prompt in refined_prompts:\n",
    "        all_prompts.append(prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "np.random.seed = 42\n",
    "randomly_sampled_examples = np.random.choice(all_prompts, 5000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "text_features_full = {}\n",
    "\n",
    "text = clip.tokenize(randomly_sampled_examples).to(device)\n",
    "\n",
    "for clip_backbone in clip_backbones:\n",
    "    encoder = models[clip_backbone][0]\n",
    "\n",
    "    indices = range(text.shape[0])\n",
    "    batches = np.array_split(indices, 10)\n",
    "\n",
    "    text_features_full[clip_backbone] = []\n",
    "\n",
    "    for batch in batches:\n",
    "        with torch.no_grad():\n",
    "            text_features_full[clip_backbone].append(encoder.encode_text(text[batch]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for clip_backbone in clip_backbones:\n",
    "    text_features_full[clip_backbone] = torch.cat(text_features_full[clip_backbone])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import clip\n",
    "import numpy as np\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "class CirclesModel(nn.Module):\n",
    "    def __init__(self, device=None, models=None, preprocesses=None, yolo_models=None, text_features_bias=None) -> None:\n",
    "        super().__init__()\n",
    "        \n",
    "        if device:\n",
    "            self.device = device\n",
    "        else:\n",
    "            self.device = \"cuda\" if torch.cuda.is_available() else \"mps\" if torch.backends.mps.is_available() else \"cpu\"\n",
    "        \n",
    "        if not models or not preprocesses:\n",
    "            raise ValueError('Models and preprocesses for CLIP model should be provided')\n",
    "\n",
    "        self.models = models\n",
    "        self.preprocesses = preprocesses\n",
    "\n",
    "        self.clip_backbones = list(self.models.keys())\n",
    "        \n",
    "        if not yolo_models:\n",
    "            raise ValueError('Models for YOLO should be provided')\n",
    "        self.yolo_models = yolo_models\n",
    "\n",
    "        if not text_features_bias:\n",
    "            self.text_features_bias = {}\n",
    "            for backbone in self.clip_backbones:\n",
    "                text_features_bias[backbone] = None\n",
    "        else:\n",
    "            self.text_features_bias = text_features_bias\n",
    "\n",
    "        self.visual_augmentation = 2\n",
    "\n",
    "        self.transform_to_tensor = T.Compose([\n",
    "            transforms.ToTensor()\n",
    "        ])\n",
    "\n",
    "    def forward(self, indices, images, prompts_list):\n",
    "        self.device = indices.device\n",
    "        if indices.is_cuda:\n",
    "            self.device_index = int(str(self.device)[-1])\n",
    "        else:\n",
    "            self.device_index = 0\n",
    "\n",
    "        # -- Getting the right data and moving it to the correct device --\n",
    "\n",
    "        # Images remain on the CPU because they are PIL Images, not Tensors\n",
    "        # Converting to Tensors leads to errors with YOLO\n",
    "        images = [images[i] for i in indices]\n",
    "\n",
    "        prompts_list = [prompts_list[i] for i in indices]\n",
    "        prompts_list = self.update_prompts_with_this_is(prompts_list)\n",
    "        prompts_tensor = [clip.tokenize(prompt_list).to(self.device) for prompt_list in prompts_list]\n",
    "\n",
    "        # -- Actual processing --\n",
    "\n",
    "        bounding_boxes = self.get_bounding_boxes(images)\n",
    "\n",
    "        # It contains the predicted bounding box for each image for each prompt\n",
    "        # Then, it is a list of length len(images) and for each entry there is a\n",
    "        # list with len(prompts[i]), where i is the i-th image \n",
    "        overall_outputs = []\n",
    "\n",
    "        with torch.no_grad():\n",
    "            for idx, prompts_tensor_for_sample in enumerate(prompts_tensor):\n",
    "                # Image crops\n",
    "                image_crops = self.get_visual_prompts(images[idx], bounding_boxes[idx])\n",
    "\n",
    "                preprocessed_image_crops = {}\n",
    "                for backbone in self.clip_backbones:\n",
    "                    preprocessed_image_crops[backbone] = torch.stack([self.preprocesses[backbone][self.device_index](image) for image in image_crops]).to(self.device)\n",
    "\n",
    "                similarities = {}\n",
    "                for backbone in self.clip_backbones:\n",
    "                    # print(len(preprocessed_image_crops[backbone]))\n",
    "                    # print(range(len(preprocessed_image_crops[backbone])))\n",
    "                    # print(preprocessed_image_crops[backbone].shape, len(preprocessed_image_crops[backbone]))\n",
    "                    pic_batches = np.array_split(range(len(preprocessed_image_crops[backbone])), len(preprocessed_image_crops[backbone]) // self.visual_augmentation)\n",
    "                    \n",
    "                    visual_features = []\n",
    "                    for pic_batch in pic_batches:\n",
    "                        visual_features.append(self.models[backbone][self.device_index].encode_image(preprocessed_image_crops[backbone][pic_batch]))\n",
    "                    visual_features = torch.cat(visual_features)\n",
    "\n",
    "                    # print(visual_features.shape)\n",
    "\n",
    "                    text_features = self.models[backbone][self.device_index].encode_text(prompts_tensor_for_sample)\n",
    "\n",
    "                    if self.text_features_bias[backbone] is not None:\n",
    "                        text_features = torch.cat([text_features, self.text_features_bias[backbone].to(self.device)])\n",
    "\n",
    "                    similarities[backbone] = cosine_similarity(visual_features, text_features)\n",
    "\n",
    "\n",
    "                similarity = torch.empty_like(similarities[self.clip_backbones[0]])\n",
    "                for prompt_idx in range(similarity.shape[0]):\n",
    "                    for proposal_idx in range(similarity.shape[1]):\n",
    "                        similarity[prompt_idx, proposal_idx] = torch.mean(torch.stack([\n",
    "                            similarities[backbone][prompt_idx, proposal_idx] for backbone in self.clip_backbones\n",
    "                        ]))\n",
    "\n",
    "\n",
    "                average = similarity.mean(dim=0)\n",
    "                scores = (similarity - average)[range(len(prompts_tensor_for_sample))]\n",
    "\n",
    "                final_scores = torch.empty((scores.shape[0], scores.shape[1] // self.visual_augmentation))\n",
    "\n",
    "                for prompt_idx in range(final_scores.shape[0]):\n",
    "                    for final_score_idx, proposal_idx in enumerate(range(0, scores.shape[1], self.visual_augmentation)):\n",
    "                        final_scores[prompt_idx, final_score_idx] = torch.max(torch.stack([\n",
    "                            scores[prompt_idx, proposal_idx + i] for i in range(self.visual_augmentation)\n",
    "                        ]))\n",
    "                # print(final_scores)\n",
    "\n",
    "                \n",
    "                _, max_indices = final_scores.max(dim=-1)\n",
    "                try:\n",
    "                    for max_idx in max_indices:\n",
    "                        overall_outputs.append(\n",
    "                            (bounding_boxes[idx][max_idx, 0:4]).to(self.device)\n",
    "                        )\n",
    "                except:\n",
    "                    for max_idx in max_indices:\n",
    "                        overall_outputs.append(\n",
    "                            torch.tensor((0, 0, 0, 0)).to(self.device)\n",
    "                        )\n",
    "\n",
    "        return torch.stack(overall_outputs)\n",
    "\n",
    "    def get_prompts(self, sample):\n",
    "        return [prompt['sent'] for prompt in sample['image']['sentences']]\n",
    "\n",
    "    def update_prompts_with_this_is(self, prompts):\n",
    "        return [['This is ' + prompt for prompt in sample] for sample in prompts]\n",
    "\n",
    "    def get_bounding_boxes(self, pil_images):\n",
    "        bounding_boxes = self.yolo_models[self.device_index].predict(pil_images, verbose=False)\n",
    "        bounding_boxes = [torch.cat([box.xyxy for box in res.boxes]) if res.boxes else None for res in bounding_boxes]\n",
    "        # bounding_boxes.show()\n",
    "        return bounding_boxes\n",
    "    \n",
    "    def draw_circle(self, image_alpha, bounding_box):\n",
    "        new_img = Image.new('RGBA', image_alpha.size, (0, 0, 0, 0))\n",
    "        draw = ImageDraw.Draw(new_img)\n",
    "        draw.ellipse((bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item()),\n",
    "                     outline='red', width=2)\n",
    "\n",
    "        new_img = Image.alpha_composite(image_alpha, new_img)\n",
    "\n",
    "        return new_img\n",
    "\n",
    "    def draw_circle_and_darken(self, image_alpha, bounding_box):\n",
    "        circle_mask = Image.new('L', image_alpha.size, 0)\n",
    "        draw = ImageDraw.Draw(circle_mask)\n",
    "        draw.ellipse((bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item()), fill=255)\n",
    "\n",
    "        alpha = 0.15\n",
    "        # alpha = 0.7\n",
    "        darkened_image = Image.new('RGB', image_alpha.size, (0, 0, 0))\n",
    "        darkened_image.paste(image_alpha.convert('RGB'), mask=circle_mask)\n",
    "        blurred_mask = circle_mask.filter(ImageFilter.GaussianBlur(radius=10))\n",
    "        darkened_image.putalpha(blurred_mask.point(lambda x: alpha * (255 - x)))\n",
    "        darkened_image = Image.alpha_composite(image_alpha, darkened_image.convert('RGBA'))\n",
    "\n",
    "        draw = ImageDraw.Draw(darkened_image)\n",
    "        draw.ellipse((bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item()),\n",
    "                     outline='red', width=4)\n",
    "\n",
    "        return darkened_image\n",
    "    \n",
    "    def draw_circle_small_and_darken(self, image_alpha, bounding_box):\n",
    "        # radius = int(0.06 * min(image_alpha.size))\n",
    "        # radius = int(min(\n",
    "        #     bounding_box[2].item() - bounding_box[0].item(),\n",
    "        #     bounding_box[3].item() - bounding_box[1].item()) / 2)\n",
    "        # center_x =  (bounding_box[0].item() + bounding_box[2].item()) / 2\n",
    "        # center_y =  (bounding_box[1].item() + bounding_box[3].item()) / 2\n",
    "\n",
    "        bbox = (bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item())\n",
    "        # bbox = (center_x - radius, center_y - radius, center_x + radius, center_y + radius)\n",
    "\n",
    "        circle_mask = Image.new('L', image_alpha.size, 0)\n",
    "        draw = ImageDraw.Draw(circle_mask)\n",
    "        draw.ellipse(bbox, fill=255)\n",
    "\n",
    "        alpha = 0.15\n",
    "        alpha = 0.4\n",
    "        # thickness = int(0.01 * min(image_alpha.size))\n",
    "        thickness = 3\n",
    "        darkened_image = Image.new('RGB', image_alpha.size, (0, 0, 0))\n",
    "        darkened_image.paste(image_alpha.convert('RGB'), mask=circle_mask)\n",
    "        blurred_mask = circle_mask.filter(ImageFilter.GaussianBlur(radius=10))\n",
    "        darkened_image.putalpha(blurred_mask.point(lambda x: alpha * (255 - x)))\n",
    "        darkened_image = Image.alpha_composite(image_alpha, darkened_image.convert('RGBA'))\n",
    "\n",
    "        draw = ImageDraw.Draw(darkened_image)\n",
    "        draw.ellipse(bbox,\n",
    "                     outline='red', width=thickness)\n",
    "\n",
    "        return darkened_image\n",
    "    \n",
    "    def draw_circle_small_and_blur_and_darken(self, image_alpha, bounding_box):\n",
    "        # radius = int(0.06 * min(image_alpha.size))\n",
    "        # radius = int(min(\n",
    "        #     bounding_box[2].item() - bounding_box[0].item(),\n",
    "        #     bounding_box[3].item() - bounding_box[1].item()) / 2)\n",
    "        # center_x =  (bounding_box[0].item() + bounding_box[2].item()) / 2\n",
    "        # center_y =  (bounding_box[1].item() + bounding_box[3].item()) / 2\n",
    "\n",
    "        bbox = (bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item())\n",
    "        # bbox = (center_x - radius, center_y - radius, center_x + radius, center_y + radius)\n",
    "\n",
    "        circle_mask = Image.new('L', image_alpha.size, 0)\n",
    "        draw = ImageDraw.Draw(circle_mask)\n",
    "        draw.ellipse(bbox, fill=255)\n",
    "\n",
    "        alpha = 0.3\n",
    "        # alpha = 0.2\n",
    "        # thickness = int(0.01 * min(image_alpha.size))\n",
    "        thickness = 2\n",
    "        darkened_image = Image.new('RGB', image_alpha.size, (0, 0, 0))\n",
    "        darkened_image.paste(image_alpha.convert('RGB'), mask=circle_mask)\n",
    "        blurred_mask = circle_mask.filter(ImageFilter.GaussianBlur(radius=10))\n",
    "        darkened_image.putalpha(blurred_mask.point(lambda x: alpha * (255 - x)))\n",
    "        darkened_image = Image.alpha_composite(image_alpha.filter(ImageFilter.GaussianBlur(radius=1.5)), darkened_image.convert('RGBA'))\n",
    "        darkened_image.paste(image_alpha.convert('RGB'), mask=circle_mask)\n",
    "\n",
    "        draw = ImageDraw.Draw(darkened_image)\n",
    "        draw.ellipse(bbox,\n",
    "                     outline='red', width=thickness)\n",
    "\n",
    "        return darkened_image\n",
    "    \n",
    "    def draw_rectangle_and_darken(self, image_alpha, bounding_box):\n",
    "        circle_mask = Image.new('L', image_alpha.size, 0)\n",
    "        draw = ImageDraw.Draw(circle_mask)\n",
    "        draw.rectangle((bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item()), fill=255)\n",
    "\n",
    "        alpha = 0.15\n",
    "        # alpha = 0.7\n",
    "        darkened_image = Image.new('RGB', image_alpha.size, (0, 0, 0))\n",
    "        darkened_image.paste(image_alpha.convert('RGB'), mask=circle_mask)\n",
    "        blurred_mask = circle_mask.filter(ImageFilter.GaussianBlur(radius=15))\n",
    "        darkened_image.putalpha(blurred_mask.point(lambda x: alpha * (255 - x)))\n",
    "        darkened_image = Image.alpha_composite(image_alpha, darkened_image.convert('RGBA'))\n",
    "\n",
    "        draw = ImageDraw.Draw(darkened_image)\n",
    "        draw.rectangle((bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item()),\n",
    "                     outline='red', width=2)\n",
    "\n",
    "        return darkened_image\n",
    "\n",
    "    def get_highlighted_bounding_boxes(self, image, bounding_boxes):\n",
    "        \"\"\"\n",
    "        Bounding boxes in the form:\n",
    "        [top left x, top left y, bottom right x, bottom right y, confidence, categoy]\n",
    "        \"\"\"\n",
    "\n",
    "        highlighted_bounding_boxes = []\n",
    "\n",
    "        image_width, image_height = image.size\n",
    "        image_alpha = image.convert('RGBA')\n",
    "        \n",
    "        for bbox_idx, bounding_box in enumerate(bounding_boxes):\n",
    "            # print(f'bbox: {bounding_box}')\n",
    "\n",
    "            # new_img = Image.new('RGBA', image.size, (0, 0, 0, 0))\n",
    "            # draw = ImageDraw.Draw(new_img)\n",
    "            # draw.ellipse((bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item()), outline='red', width=4)\n",
    "\n",
    "            # new_img = Image.alpha_composite(image_alpha, new_img)\n",
    "\n",
    "            # new_img = self.draw_circle(image_alpha, bounding_box)\n",
    "            # new_img = self.draw_circle_and_darken(image_alpha, bounding_box)\n",
    "            # new_img = self.draw_circle_small_and_darken(image_alpha, bounding_box)\n",
    "            new_img = self.draw_circle_small_and_blur_and_darken(image_alpha, bounding_box)\n",
    "            # new_img = self.draw_rectangle_and_darken(image_alpha, bounding_box)\n",
    "\n",
    "            # fig, ax = plt.subplots()\n",
    "            # ax.imshow(new_img)\n",
    "            # ax.set_title(crop_centroid_normalized)\n",
    "\n",
    "\n",
    "            highlighted_bounding_boxes.append(new_img)\n",
    "\n",
    "        if len(highlighted_bounding_boxes) == 0:\n",
    "            highlighted_bounding_boxes.append(image)\n",
    "                \n",
    "        return highlighted_bounding_boxes\n",
    "\n",
    "    def get_image_with_marker(self, image, bbox, stroke_color='red', stroke_width=1):\n",
    "        result = image.copy()\n",
    "        draw = ImageDraw.Draw(result)\n",
    "        draw.ellipse(bbox, outline=stroke_color, width=stroke_width)\n",
    "        \n",
    "        return result\n",
    "    \n",
    "    def get_image_with_marker_and_blur(self, image, bbox, stroke_color='red', stroke_width=1, blur_radius=1):\n",
    "        result = image.filter(ImageFilter.GaussianBlur(radius=blur_radius))\n",
    "        mask = Image.new('L', image.size, 0)\n",
    "        draw = ImageDraw.Draw(mask)\n",
    "        draw.rectangle(bbox, fill=255)\n",
    "        result.paste(image, mask=mask)\n",
    "        draw = ImageDraw.Draw(result)\n",
    "        draw.ellipse(bbox, outline=stroke_color, width=stroke_width)\n",
    "        \n",
    "        return result\n",
    "\n",
    "    def get_image_with_marker_and_grayscale(self, image, bbox, stroke_color='red', stroke_width=1):\n",
    "        result = image.convert('L').convert('RGB')\n",
    "        mask = Image.new('L', image.size, 0)\n",
    "        draw = ImageDraw.Draw(mask)\n",
    "        draw.rectangle(bbox, fill=255)\n",
    "        result.paste(image, mask=mask)\n",
    "        draw = ImageDraw.Draw(result)\n",
    "        draw.ellipse(bbox, outline=stroke_color, width=stroke_width)\n",
    "\n",
    "        return result\n",
    "\n",
    "    def get_image_with_marker_and_blur_grayscale(self, image, bbox, stroke_color='red', stroke_width=1, blur_radius=1):\n",
    "        result = image.filter(ImageFilter.GaussianBlur(radius=blur_radius)).convert('L').convert('RGB')\n",
    "        mask = Image.new('L', image.size, 0)\n",
    "        draw = ImageDraw.Draw(mask)\n",
    "        draw.rectangle(bbox, fill=255)\n",
    "        result.paste(image, mask=mask)\n",
    "        draw = ImageDraw.Draw(result)\n",
    "        draw.ellipse(bbox, outline=stroke_color, width=stroke_width)\n",
    "\n",
    "        return result\n",
    "\n",
    "    def get_visual_prompts(self, image, bounding_boxes):\n",
    "        self.visual_augmentation = 2\n",
    "        visual_prompts = []\n",
    "\n",
    "        # if bounding_boxes is None:\n",
    "        #     return [image, image]\n",
    "\n",
    "        for bounding_box in (bounding_boxes if bounding_boxes is not None else []):\n",
    "            bounding_box = (bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item())\n",
    "\n",
    "            # For the following line to work correctly bounding boxes should actually be removed from\n",
    "            # YOLO's results, as that's what is actually used in the end\n",
    "            # if (bounding_box[2] - bounding_box[0]) * (bounding_box[3] - bounding_box[1]) >= (image.size[0] * image.size[1]) * 0.8:\n",
    "            #     continue\n",
    "\n",
    "            stroke_color = 'red'\n",
    "            stroke_width = 3\n",
    "            blur_radius = 20\n",
    "\n",
    "            # fig, ax = plt.subplots()\n",
    "            # ax.imshow(self.get_image_with_marker(image, bounding_box, stroke_color=stroke_color, stroke_width=stroke_width))\n",
    "            # fig, ax = plt.subplots()\n",
    "            # ax.imshow(self.get_image_with_marker_and_blur(image, bounding_box, stroke_color=stroke_color, stroke_width=stroke_width, blur_radius=blur_radius))\n",
    "\n",
    "            bbox_visual_prompts = [\n",
    "                # self.get_image_with_marker(image, bounding_box, stroke_color=stroke_color, stroke_width=stroke_width),\n",
    "                self.get_image_with_marker_and_blur(image, bounding_box, stroke_color=stroke_color, stroke_width=stroke_width, blur_radius=blur_radius),\n",
    "                # self.get_image_with_marker_and_grayscale(image, bounding_box, stroke_color=stroke_color, stroke_width=stroke_width),\n",
    "                self.get_image_with_marker_and_blur_grayscale(image, bounding_box, stroke_color=stroke_color, stroke_width=stroke_width, blur_radius=blur_radius),\n",
    "            ]\n",
    "\n",
    "\n",
    "\n",
    "            for el in bbox_visual_prompts:\n",
    "                visual_prompts.append(el)\n",
    "\n",
    "                # fig, ax = plt.subplots()\n",
    "                # ax.imshow(el)\n",
    "\n",
    "        if len(visual_prompts) == 0:\n",
    "            # If no region proposal, return the whole image.\n",
    "            # It is inserted as many times as each region would\n",
    "            # be augmented to ensure consistency in the algorithm\n",
    "            for _ in range(self.visual_augmentation):\n",
    "                visual_prompts.append(image)\n",
    "                \n",
    "        return visual_prompts\n",
    "\n",
    "\n",
    "    def get_cropped_bounding_boxes(self, image, bounding_boxes):\n",
    "        \"\"\"\n",
    "        Bounding boxes in the form:\n",
    "        [top left x, top left y, bottom right x, bottom right y, confidence, categoy]\n",
    "        \"\"\"\n",
    "\n",
    "        cropped_bounding_boxes = []\n",
    "\n",
    "        image_width, image_height = image.size\n",
    "        \n",
    "        for bbox_idx, bounding_box in enumerate(bounding_boxes):\n",
    "            # print(f'bbox: {bounding_box}')\n",
    "\n",
    "            cropped_img = image.crop((bounding_box[0].item(), bounding_box[1].item(), bounding_box[2].item(), bounding_box[3].item()))\n",
    "\n",
    "            # cropped_img\n",
    "\n",
    "            # Centroid: (min + (max - min) / 2) / dimension\n",
    "            crop_centroid_normalized = (\n",
    "                (bounding_box[0].item() + (bounding_box[2].item() - bounding_box[0].item()) / 2) / image_width,\n",
    "                (bounding_box[1].item() + (bounding_box[3].item() - bounding_box[1].item()) / 2 ) / image_height\n",
    "            )\n",
    "\n",
    "            if crop_centroid_normalized[0] < 0.5:\n",
    "                overlay = Image.new('RGBA', cropped_img.size, overlay_colors[0])\n",
    "            elif crop_centroid_normalized[0] > 0.5:\n",
    "                overlay = Image.new('RGBA', cropped_img.size, overlay_colors[1])\n",
    "            else:\n",
    "                overlay = Image.new('RGBA', cropped_img.size, overlay_colors[-1])\n",
    "            blended = Image.alpha_composite(cropped_img.convert('RGBA'), overlay)\n",
    "            cropped_bounding_boxes.append(blended)\n",
    "\n",
    "            # blended.show()\n",
    "            # fig, ax = plt.subplots()\n",
    "            # ax.imshow(blended)\n",
    "            # ax.set_title(crop_centroid_normalized)\n",
    "\n",
    "\n",
    "            # cropped_bounding_boxes.append(cropped_img)\n",
    "\n",
    "        if len(cropped_bounding_boxes) == 0:\n",
    "            cropped_bounding_boxes.append(image)\n",
    "                \n",
    "        return cropped_bounding_boxes\n",
    "\n",
    "circles_model = CirclesModel(models=models, preprocesses=preprocesses, yolo_models=yolo_models, text_features_bias=text_features_full)\n",
    "\n",
    "overlay_colors = [\n",
    "    # (0, 0, 0, 0),       # None,\n",
    "    # (0, 0, 0, 0),       # None\n",
    "    (255, 0, 0, 128),   # Red, alpha = 0.5\n",
    "    (0, 255, 0, 128),   # Green, alpha = 0.5\n",
    "    (0, 0, 255, 128),   # Blue, alpha = 0.5\n",
    "    (0, 0, 0, 0),       # None\n",
    "]\n",
    "\n",
    "if torch.cuda.device_count() > 1:\n",
    "    circles_model = torch.nn.DataParallel(circles_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.ops import box_iou\n",
    "\n",
    "def iou_metric(bounding_boxes, ground_truth_bounding_boxes):\n",
    "    \"\"\"\n",
    "    Localization Accuracy Metric\n",
    "\n",
    "    Intersection over Union (IoU) is a common metric measure for localization accuracy.\n",
    "    \"\"\"\n",
    "\n",
    "    ground_truth_bounding_boxes = torch.tensor(ground_truth_bounding_boxes).unsqueeze(0).to(device)\n",
    "\n",
    "    return box_iou(bounding_boxes, ground_truth_bounding_boxes)\n",
    "\n",
    "def cosine_similarity_metric(bounding_boxes, ground_truth_bounding_boxes):\n",
    "    \"\"\"\n",
    "    Cosine Similarity Metric\n",
    "\n",
    "    Cosine similarity is a common metric measure for semantic similarity.\n",
    "    \"\"\"\n",
    "\n",
    "    ground_truth_bounding_boxes = torch.tensor(ground_truth_bounding_boxes).to(device)\n",
    "    \n",
    "    return cosine_similarity(bounding_boxes, ground_truth_bounding_boxes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "idx = 17 # man on the beach with frisbee\n",
    "# idx = 20 # motorbikes\n",
    "# idx = 22 # cows on a beach\n",
    "# idx = 25 # three oranges and a banana\n",
    "# idx = 32 # guy with a horse and two busses\n",
    "# idx = 33 # luggage\n",
    "# idx = 34 # man on bed in front of a window\n",
    "idx = 35 # two zebras\n",
    "# idx = 36 # two horses\n",
    "# idx = 38 # chairs around a table with some sweets on top\n",
    "# idx = 39 # two monitors\n",
    "# idx = 42 # folks playing wii\n",
    "# idx = 43 # yellow vehicle and surfboard\n",
    "# idx = 44 # two women playing tennis\n",
    "# idx = 45 # woman with a thing of bananas\n",
    "# idx = 46 # industrial kitchen stove\n",
    "# idx = 47 # two guys, one has a beard\n",
    "# idx = 49 # girl eating pizza\n",
    "# idx = 50 # vertical fork\n",
    "# idx = 51 # sandwiches\n",
    "# idx = 54 # computer on the right\n",
    "# idx = 57 # woman playing tennis\n",
    "\n",
    "img = next(iter(test_loader))[0][idx]\n",
    "bbox_gt = next(iter(test_loader))[1][idx]\n",
    "prompt = next(iter(test_loader))[2][idx]\n",
    "\n",
    "img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt, bbox_gt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# prompt = ['the man on the right with a red overlay', 'the man with a blue shirt']#, 'a photo of a man who is about to throw a frisbee'] # idx == 17\n",
    "\n",
    "# prompt = ['the red motorcycle with a blue overlay'] # idx == 20\n",
    "# prompt = ['a red & black color bike in ftont of the three guys'] # idx == 20\n",
    "\n",
    "# prompt = ['the smaller animal'] # idx == 22\n",
    "\n",
    "# prompt = ['the orange closest to the banana',\n",
    "#     'orange with a green overlay',\n",
    "#  'orange between other oranges and a banana',\n",
    "    # 'A photo of a orange',\n",
    "    # 'A photo of a dining table',\n",
    "    # 'A photo of a banana'\n",
    "    # ] # idx == 25\n",
    "\n",
    "# prompt = ['the orange closest to the banana with a red overlay']\n",
    "\n",
    "# prompt = ['near zebra with a red overlay', 'zebra eating grass with a red overlay'] # idx == 35\n",
    "\n",
    "# prompt = ['the man with glasses'] # idx == 32\n",
    "\n",
    "# prompt = [\n",
    "#     'a man with beard wearing blue shirt with his friend',\n",
    "#     'a man with a beard',\n",
    "# ] # idx == 47\n",
    "\n",
    "# prompt = ['the right computer in the right hand picture with a green overlay',\n",
    "#  'the computer on the right in the right hand picture with a green overlay'] # idx == 54\n",
    "\n",
    "# prompt = ['the woman on the right']# with a green overlay',\n",
    "#  'the girl with the racket in the photo on the right with a green overlay'] # idx == 57"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# transform = T.Compose([\n",
    "#     T.Resize(size=224, interpolation=T.InterpolationMode.BICUBIC, max_size=None, antialias='warn'),\n",
    "#     T.CenterCrop(size=(224, 224)),\n",
    "#     T.ToTensor(),\n",
    "# ])\n",
    "\n",
    "# img_tensor = transform(img)\n",
    "# print(img_tensor.shape)\n",
    "# res = yolo_models[0](torch.stack([img_tensor]))\n",
    "\n",
    "res = yolo_models[0](img)\n",
    "# res.pred[0].cpu().numpy()[:, -1]\n",
    "res.pred[0].cpu().numpy() "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Draw a circle onto the image\n",
    "from PIL import Image, ImageDraw, ImageFilter\n",
    "\n",
    "for res_pred in res.pred[0]:\n",
    "    bbox = res_pred[0:4].cpu().numpy()\n",
    "\n",
    "    circle_mask = Image.new('L', img.size, 0)\n",
    "    draw = ImageDraw.Draw(circle_mask)\n",
    "    draw.ellipse(bbox, fill=255)\n",
    "\n",
    "    alpha = 0.2\n",
    "    darkened_image = Image.new('RGB', img.size, (0, 0, 0))\n",
    "    darkened_image.paste(img, mask=circle_mask)\n",
    "    blurred_mask = circle_mask.filter(ImageFilter.GaussianBlur(radius=10))\n",
    "    darkened_image.putalpha(blurred_mask.point(lambda x: alpha * (255 - x)))\n",
    "    darkened_image = Image.alpha_composite(img.convert('RGBA').filter(ImageFilter.GaussianBlur(radius=2)), darkened_image.convert('RGBA'))\n",
    "\n",
    "    darkened_image.paste(img, mask=circle_mask)\n",
    "\n",
    "    draw = ImageDraw.Draw(darkened_image)\n",
    "    draw.ellipse(bbox, outline='red', width=4)\n",
    "\n",
    "    # img_new = Image.new('RGBA', img.size, (0, 0, 0, 0))\n",
    "    # draw = ImageDraw.Draw(img_new)\n",
    "    # draw.ellipse(bbox, outline='red', width=4)\n",
    "\n",
    "    # img_new = Image.alpha_composite(img.convert('RGBA'), img_new)\n",
    "\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.imshow(darkened_image)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "indices = torch.tensor([range(1)]).to(device)\n",
    "outputs = circles_model(indices, [img], [prompt])#['the woman on the left, the girl with the racket in the photo on the left'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "output_idx = 0\n",
    "\n",
    "# Loading the image\n",
    "# img = img[output_idx]\n",
    "\n",
    "# Preparing the output\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "# Display the image\n",
    "ax.imshow(img)\n",
    "\n",
    "colors = ['r', 'b', 'g', 'yellow', 'orange']\n",
    "\n",
    "# print(outputs.shape, bbox.unsqueeze(0).shape)\n",
    "\n",
    "gt_and_outputs = torch.cat([torch.tensor(bbox_gt).unsqueeze(0).to(device), outputs])\n",
    "# print(gt_and_outputs, bbox_gt)\n",
    "\n",
    "# Create a Rectangle patch\n",
    "for bbox_idx, (bbox, color) in enumerate(zip(gt_and_outputs, colors)):\n",
    "    bounding_box_coordinates = bbox.cpu()\n",
    "    top_left_x, top_left_y = bounding_box_coordinates[0], bounding_box_coordinates[1]\n",
    "    width, height = bounding_box_coordinates[2]- top_left_x, bounding_box_coordinates[3] - top_left_y\n",
    "\n",
    "    # Parameters: (x, y), width, height\n",
    "    rect = patches.Rectangle((top_left_x, top_left_y), width, height, linewidth=2 if bbox_idx == 0 else 1, edgecolor=color, facecolor='none')\n",
    "\n",
    "    # Add the patch to the Axes\n",
    "    ax.add_patch(rect)\n",
    "\n",
    "ax.set_title(prompt)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### To compute average cosine similarity between embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images, gt_bounding_boxes, prompts = next(iter(test_loader))\n",
    "\n",
    "idx_to_use = 3\n",
    "\n",
    "indices = torch.tensor(list(range(len(images)))).to(device)\n",
    "outputs = circles_model(indices[0:2], images[idx_to_use:idx_to_use+2], prompts[idx_to_use:idx_to_use+2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax.imshow(images[idx_to_use].crop())\n",
    "\n",
    "for output in outputs:\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.imshow(images[idx_to_use].crop(output.cpu().numpy()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "overall_outputs = []\n",
    "\n",
    "for batch_idx, (images, gt_bounding_boxes, prompts) in enumerate(test_loader):\n",
    "    print(f'-- Batch index: {batch_idx} --')\n",
    "\n",
    "    prompts_tensor = [clip.tokenize(prompt_list) for prompt_list in prompts]\n",
    "    \n",
    "    indices = torch.tensor(list(range(len(images)))).to(device)\n",
    "    outputs = circles_model(indices, images, prompts)\n",
    "\n",
    "    break\n",
    "\n",
    "    overall_outputs.append(outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cos_sim_cpu = []\n",
    "for out in overall_outputs:\n",
    "    for cos_sim_val in out:\n",
    "        cos_sim_cpu.append(cos_sim_val.item())\n",
    "cos_sim_cpu = np.array(cos_sim_cpu)\n",
    "np.nanmean(cos_sim_cpu)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### To compute standard metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from torchvision.ops import boxes as box_ops\n",
    "\n",
    "IoUs = []\n",
    "cosine_similarities = []\n",
    "  \n",
    "for batch_idx, (images, gt_bounding_boxes, prompts) in enumerate(test_loader):\n",
    "    print(f'-- Batch index: {batch_idx} --')\n",
    "\n",
    "    # if batch_idx < 62:\n",
    "    #     continue\n",
    "\n",
    "    prompts_tensor = [clip.tokenize(prompt_list) for prompt_list in prompts]\n",
    "    \n",
    "    indices = torch.tensor(list(range(len(images)))).to(device)\n",
    "    outputs = circles_model(indices, images, prompts)\n",
    "    # outputs = circles_model(indices[0:1], images[21:22], prompts[21:22])\n",
    "    # outputs = circles_model(indices[0:1], images[22:23], prompts[22:23])\n",
    "    \n",
    "    outputs_grouped_by_sample = []\n",
    "    outputs_idx = 0\n",
    "    prompts_idx = 0\n",
    "    while True:\n",
    "        if not prompts_idx < len(images):\n",
    "            break\n",
    "\n",
    "        outputs_grouped_by_sample.append(\n",
    "            outputs[outputs_idx : outputs_idx + len(prompts[prompts_idx])]\n",
    "        )\n",
    "\n",
    "        outputs_idx += len(prompts[prompts_idx])\n",
    "        prompts_idx += 1\n",
    "\n",
    "    for output_bboxes, gt_bboxes in zip(outputs_grouped_by_sample, gt_bounding_boxes):\n",
    "        \"\"\"\n",
    "        There is one output bounding box for each prompt given in input.\n",
    "        Note that each prompt for a given input is actually a list of prompts,\n",
    "        therefore it can contain an arbitrary number of promps. Hence, there is\n",
    "        a bounding box for each one of them.\n",
    "        \"\"\"\n",
    "\n",
    "        result_ious = iou_metric(output_bboxes, gt_bboxes)\n",
    "        result_cosine_similarity = cosine_similarity_metric(output_bboxes, gt_bboxes)\n",
    "\n",
    "        for iou in result_ious:\n",
    "            IoUs.append(iou)\n",
    "\n",
    "        for cs in result_cosine_similarity:\n",
    "            cosine_similarities.append(cs)\n",
    "\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "counter = 0\n",
    "counter_threshold = 0\n",
    "for iou in IoUs:\n",
    "    if iou == 0:\n",
    "        counter += 1\n",
    "    # if iou < 0.5:\n",
    "    #     counter_threshold += 1\n",
    "    if iou >= 0.5:\n",
    "        counter_threshold += 1\n",
    "counter, counter_threshold, len(IoUs), counter_threshold / len(IoUs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IoUs_to_cpu = np.array([tensor.item() if torch.is_tensor(tensor) else 0 for tensor in IoUs])\n",
    "mIoU = np.nanmean(IoUs_to_cpu)\n",
    "\n",
    "cosine_similarities_to_cpu = np.array([tensor.item() if torch.is_tensor(tensor) else 0 for tensor in cosine_similarities])\n",
    "m_cos_sim = np.nanmean(cosine_similarities_to_cpu)\n",
    "\n",
    "print('--- Metrics ---')\n",
    "print(f'Mean Intersection over Union (mIoU): {mIoU}')\n",
    "print(f'Mean Cosine Similarity: {m_cos_sim}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "np.savetxt('outcomes/iou_using_2_visual_augmentation_only_ViT-L-14-336-YOLOv8.csv', IoUs_to_cpu, delimiter=',')\n",
    "np.savetxt('outcomes/cossim_using_2_visual_augmentation_only_ViT-L-14-336-YOLOv8.csv', cosine_similarities_to_cpu, delimiter=',')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "output_idx = 0\n",
    "\n",
    "# Loading the image\n",
    "img = images[output_idx]\n",
    "\n",
    "# Preparing the output\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "# Display the image\n",
    "ax.imshow(img)\n",
    "\n",
    "colors = ['r', 'b', 'g']\n",
    "\n",
    "# Create a Rectangle patch\n",
    "for bbox, color in zip(outputs_grouped_by_sample[output_idx][1:2], colors):\n",
    "    bounding_box_coordinates = bbox.cpu()\n",
    "    top_left_x, top_left_y = bounding_box_coordinates[0], bounding_box_coordinates[1]\n",
    "    width, height = bounding_box_coordinates[2]- top_left_x, bounding_box_coordinates[3] - top_left_y\n",
    "\n",
    "    # Parameters: (x, y), width, height\n",
    "    rect = patches.Rectangle((top_left_x, top_left_y), width, height, linewidth=1, edgecolor=color, facecolor='none')\n",
    "\n",
    "    # Add the patch to the Axes\n",
    "    ax.add_patch(rect)\n",
    "\n",
    "ax.set_title(prompts[output_idx][1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "available_gpus = [torch.cuda.device(i) for i in range(torch.cuda.device_count())]\n",
    "available_gpus"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "venv",
   "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.11"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}