text_Segmentation.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install opencv-python matplotlib\n",
    "# !pip install 'git+https://github.com/facebookresearch/segment-anything-2.git'\n",
    "\n",
    "# !mkdir -p images\n",
    "# !wget -P images https://raw.githubusercontent.com/facebookresearch/segment-anything-2/main/notebooks/images/cars.jpg\n",
    "\n",
    "# !mkdir -p ./checkpoints/\n",
    "# !wget -P ./checkpoints/ https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_large.pt\n",
    "# !wget -P ./checkpoints/ https://dl.fbaipublicfiles.com/segment_anything_2/072824/sam2_hiera_small.pt\n",
    "\n",
    "# INSTALL CLIP\n",
    "# !pip install ftfy regex tqdm\n",
    "# !pip install git+https://github.com/pablogvrcia/CLIP.git\n",
    "\n",
    "# INSTALL GOOGLE GENAI\n",
    "# !pip install -U -q google-generativeai\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/pablo/miniconda3/envs/clip/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/home/pablo/miniconda3/envs/clip/lib/python3.12/site-packages/sam2/modeling/sam/transformer.py:23: UserWarning: Flash Attention is disabled as it requires a GPU with Ampere (8.0) CUDA capability.\n",
      "  OLD_GPU, USE_FLASH_ATTN, MATH_KERNEL_ON = get_sdpa_settings()\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import torch\n",
    "import torchvision\n",
    "from PIL import Image\n",
    "import CLIP.clip as clip\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import google.generativeai as genai\n",
    "import time\n",
    "from sam2.build_sam import build_sam2\n",
    "from sam2.automatic_mask_generator import SAM2AutomaticMaskGenerator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "PyTorch version: 2.4.1+cu121\n",
      "Torchvision version: 0.19.1+cu121\n",
      "CUDA is available: True\n",
      "CLIP functions: ['__builtins__', '__cached__', '__doc__', '__file__', '__loader__', '__name__', '__package__', '__path__', '__spec__', 'available_models', 'clip', 'load', 'model', 'simple_tokenizer', 'tokenize']\n"
     ]
    }
   ],
   "source": [
    "print(\"PyTorch version:\", torch.__version__)\n",
    "print(\"Torchvision version:\", torchvision.__version__)\n",
    "print(\"CUDA is available:\", torch.cuda.is_available())\n",
    "print(\"CLIP functions:\", dir(clip))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "\n",
    "def load_clip_model():\n",
    "    \"\"\"\n",
    "    Loads the CLIP model.\n",
    "\n",
    "    Returns:\n",
    "      The CLIP model and the preprocess function.\n",
    "    \"\"\"\n",
    "    clip_model, preprocess = clip.load(\"ViT-B/32\", device=device)\n",
    "    return clip_model, preprocess\n",
    "\n",
    "\n",
    "def load_sam_model():\n",
    "    \"\"\"\n",
    "    Loads the SAM 2 model.\n",
    "\n",
    "    Returns:\n",
    "      The SAM 2 model and the mask generator.\n",
    "    \"\"\"\n",
    "    sam2_checkpoint = \"./checkpoints/sam2_hiera_small.pt\"\n",
    "    model_cfg = \"sam2_hiera_s.yaml\"\n",
    "    sam2 = build_sam2(\n",
    "        model_cfg, sam2_checkpoint, device=device, apply_postprocessing=False\n",
    "    )\n",
    "    mask_generator = SAM2AutomaticMaskGenerator(sam2)\n",
    "    return sam2, mask_generator\n",
    "\n",
    "\n",
    "def load_gemini_model():\n",
    "    \"\"\"\n",
    "    Loads the Gemini model.\n",
    "\n",
    "    Returns:\n",
    "      The Gemini model.\n",
    "    \"\"\"\n",
    "    GEMINI_API_KEY = \"AIzaSyAVi-LojQT7143OpXmdogGjAe6yBlwWHSI\"\n",
    "\n",
    "    # Configure Gemini\n",
    "    genai.configure(api_key=GEMINI_API_KEY)\n",
    "\n",
    "    gemini_model = genai.GenerativeModel(\n",
    "        model_name=\"gemini-1.5-pro\",\n",
    "        # Create the Gemini model\n",
    "        generation_config={\n",
    "            \"temperature\": 0.1,  # Adjust temperature for more focused responses\n",
    "            \"top_p\": 0.95,\n",
    "            \"top_k\": 64,\n",
    "            \"max_output_tokens\": 8192,\n",
    "            \"response_mime_type\": \"application/json\",\n",
    "        },\n",
    "    )\n",
    "    \n",
    "    [context, example_cases] = load_gemini_context()\n",
    "    \n",
    "    return gemini_model, context, example_cases\n",
    "\n",
    "\n",
    "def load_gemini_context():\n",
    "    files = [\n",
    "        upload_to_gemini(\"./images/cars.jpg\"),\n",
    "        upload_to_gemini(\"./images/red_car.png\"),\n",
    "        upload_to_gemini(\"./images/pink_car_0.png\"),\n",
    "        upload_to_gemini(\"./images/pink_car_1.png\"),\n",
    "        upload_to_gemini(\"./images/red_car_wheel.png\"),\n",
    "        upload_to_gemini(\"./images/road.png\"),\n",
    "    ]\n",
    "\n",
    "    return [\n",
    "        [\n",
    "            \"Having this image as a reference, I will send you highlited spots in the image an you will answer with only YES or NO \",\n",
    "            files[0],\n",
    "            files[1],\n",
    "            \"Is this a part of a red car?\",\n",
    "            \"output: YES or NO YES\",\n",
    "            \"Having this image as a reference, I will send you highlited spots in the image an you will answer with only YES or NO \",\n",
    "            files[0],\n",
    "            files[1],\n",
    "            \"Is this a part of a red car?\",\n",
    "            \"output: YES or NO NO\",\n",
    "            \"Having this image as a reference, I will send you highlited spots in the image an you will answer with only YES or NO \",\n",
    "            files[0],\n",
    "            files[1],\n",
    "            \"Is this a part of a pink car?\",\n",
    "            \"output: YES or NO YES\",\n",
    "            # \"Having this image as a reference, I will send you highlited spots in the image an you will answer with only YES or NO \",\n",
    "            # files[0],\n",
    "            # files[5],\n",
    "            # \"Is this a part of a road?\",\n",
    "            # \"output: YES or NO YES\",\n",
    "            # \"Having this image as a reference, I will send you highlited spots in the image an you will answer with only YES or NO \",\n",
    "            # files[0],\n",
    "            # files[4],\n",
    "            # \"Is this a part of a road?\",\n",
    "            # \"output: YES or NO NO\",\n",
    "            # \"Having this image as a reference, I will send you highlited spots in the image an you will answer with only YES or NO \",\n",
    "            # files[0],\n",
    "            # files[4],\n",
    "            # \"Is this a part of a red car?\",\n",
    "            # \"output: YES or NO YES\",\n",
    "        ],\n",
    "        3,\n",
    "    ]\n",
    "\n",
    "\n",
    "def append_to_gemini_context(context, original, masked, text_prompt):\n",
    "    # Upload the masked image to Gemini\n",
    "    new_context = [\n",
    "        \"Having this image as a reference, I will send you highlited spots in the image an you will answer with only YES or NO \",\n",
    "        original,\n",
    "        masked,\n",
    "        f\"Is this a part of {text_prompt}?\",\n",
    "        \"output: YES or NO \",\n",
    "    ]\n",
    "\n",
    "    all_context = context + new_context\n",
    "\n",
    "    return all_context\n",
    "\n",
    "def upload_to_gemini(path, mime_type=None):\n",
    "    \"\"\"Uploads the given file to Gemini.\n",
    "\n",
    "    See https://ai.google.dev/gemini-api/docs/prompting_with_media\n",
    "    \"\"\"\n",
    "    file = genai.upload_file(path, mime_type=mime_type)\n",
    "    print(f\"Uploaded file '{file.display_name}' as: {file.uri}\")\n",
    "    return file\n",
    "\n",
    "\n",
    "def upload_masked_to_gemini(image_path, mask, mime_type=\"image/png\"):\n",
    "    \"\"\"\n",
    "    Applies the mask to the image and uploads the result to Gemini.\n",
    "    \"\"\"\n",
    "    # Apply the mask to the image\n",
    "    image = Image.open(image_path)\n",
    "    masked_image = image * np.array(mask)[..., None]\n",
    "    masked_image = Image.fromarray(masked_image.astype(\"uint8\"))\n",
    "\n",
    "    # Save the masked image temporarily\n",
    "    temp_file_path = \"temp_masked_image.png\"\n",
    "    masked_image.save(temp_file_path)\n",
    "\n",
    "    # plot masked image\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(masked_image)\n",
    "    plt.axis(\"off\")\n",
    "    plt.show()\n",
    "\n",
    "    # Upload the masked image to Gemini\n",
    "    file = upload_to_gemini(temp_file_path, mime_type=mime_type)\n",
    "\n",
    "    # Remove the temporary file\n",
    "    os.remove(temp_file_path)\n",
    "\n",
    "    return file\n",
    "\n",
    "\n",
    "def calculate_image_clip_embedding(image_path):\n",
    "    \"\"\"\n",
    "    Calculates the CLIP embedding for the given image.\n",
    "\n",
    "    Args:\n",
    "      image_path: Path to the image.\n",
    "\n",
    "    Returns:\n",
    "      A tensor of the CLIP embedding for the image.\n",
    "    \"\"\"\n",
    "    image = Image.open(image_path)\n",
    "    \n",
    "    image_array = preprocess(image).unsqueeze(0).to(device)\n",
    "    \n",
    "    with torch.no_grad():\n",
    "        image_features = clip_model.encode_image(image_array)\n",
    "    return image_features, image_array\n",
    "\n",
    "\n",
    "def calculate_text_clip_embedding(text):\n",
    "    \"\"\"\n",
    "    Calculates the CLIP embedding for the given text.\n",
    "\n",
    "    Args:\n",
    "      text: Text to calculate the embedding for.\n",
    "\n",
    "    Returns:\n",
    "      A tensor of the CLIP embedding for the text.\n",
    "    \"\"\"\n",
    "    text_token = clip.tokenize([text]).to(device)\n",
    "    with torch.no_grad():\n",
    "        text_features = clip_model.encode_text(text_token)\n",
    "    return text_features\n",
    "\n",
    "\n",
    "def calculate_clips_similarity(image_path, generated_masks, text_prompt):\n",
    "    \"\"\"\n",
    "    Calculates the similarity between the text prompt and the generated masks.\n",
    "\n",
    "    Args:\n",
    "      image_path: Path to the image.\n",
    "      generated_masks: A list of masks generated by SAM 2.\n",
    "      text_prompt: Text prompt describing the object to segment.\n",
    "\n",
    "    Returns:\n",
    "      A list of similarities between the text and mask features.\n",
    "    \"\"\"\n",
    "    # Load and preprocess the image\n",
    "    image = Image.open(image_path)\n",
    "\n",
    "    # Generate CLIP embedding for the text prompt\n",
    "    text_features = calculate_text_clip_embedding(text_prompt)\n",
    "\n",
    "    # Calculate CLIP embeddings for the candidate masks\n",
    "    mask_features = calculate_clip_embeddings_for_masks(\n",
    "        generated_masks, np.array(image)\n",
    "    )\n",
    "\n",
    "    # Normalize the features and calculate the similarity\n",
    "    text_features = text_features / text_features.norm(dim=-1, keepdim=True)\n",
    "    mask_features = mask_features / mask_features.norm(dim=-1, keepdim=True)\n",
    "    similarity = (100.0 * text_features @ mask_features.T).softmax(dim=-1)\n",
    "    return similarity[0].tolist()\n",
    "\n",
    "\n",
    "def calculate_sam_masks(image_path):\n",
    "    \"\"\"\n",
    "    Calculates masks using SAM 2 for the given image.\n",
    "\n",
    "    Args:\n",
    "      image_path: Path to the image.\n",
    "\n",
    "    Returns:\n",
    "      A list of masks generated by SAM 2.\n",
    "    \"\"\"\n",
    "    image = Image.open(image_path)\n",
    "    image_array = np.array(image)\n",
    "    masks = mask_generator.generate(image_array)\n",
    "\n",
    "    # Extract masks and scores (adapt this based on SAM 2 output format)\n",
    "    generated_masks = [mask[\"segmentation\"] for mask in masks]\n",
    "    scores = [mask[\"stability_score\"] for mask in masks]\n",
    "\n",
    "    return [generated_masks, scores]\n",
    "\n",
    "\n",
    "def calculate_clip_embeddings_for_masks(masks, image):\n",
    "    \"\"\"\n",
    "    Calculates the CLIP embeddings for the given masks.\n",
    "\n",
    "    Args:\n",
    "      masks: A list of masks.\n",
    "\n",
    "    Returns:\n",
    "      A tensor of the CLIP embeddings for the masks.\n",
    "    \"\"\"\n",
    "    mask_features = []\n",
    "    for mask in masks:\n",
    "        masked_image = image * np.array(mask)[..., None]\n",
    "        masked_image_array = (\n",
    "            preprocess(Image.fromarray(masked_image.astype(\"uint8\")))\n",
    "            .unsqueeze(0)\n",
    "            .to(device)\n",
    "        )\n",
    "        with torch.no_grad():\n",
    "            mask_feature = clip_model.encode_image(masked_image_array)\n",
    "        mask_features.append(mask_feature)\n",
    "    mask_features = torch.cat(mask_features, dim=0)\n",
    "    return mask_features\n",
    "\n",
    "\n",
    "def calculate_gemini_response(image_path, generated_masks, text_prompt):\n",
    "    \"\"\"\n",
    "    Calculates the response from Gemini for the given image and masks.\n",
    "\n",
    "    Args:\n",
    "      image_path: Path to the image.\n",
    "      generated_masks: A list of masks generated by SAM 2.\n",
    "      text_prompt: Text prompt describing the object to segment.\n",
    "\n",
    "    Returns:\n",
    "      The response from Gemini.\n",
    "    \"\"\"\n",
    "    original_file = upload_to_gemini(image_path)\n",
    "\n",
    "    gemini_responses = []\n",
    "    counter = 0\n",
    "    for mask in generated_masks:\n",
    "        # just for testing purposes\n",
    "        if counter < 0:\n",
    "            # Upload masked image to Gemini\n",
    "            masked_file = upload_masked_to_gemini(image_path, mask)\n",
    "\n",
    "            # Prepare the prompt for Gemini\n",
    "            new_context = append_to_gemini_context(\n",
    "                context, original_file, masked_file, text_prompt\n",
    "            )\n",
    "\n",
    "            # Get Gemini's response\n",
    "            try:\n",
    "                response = gemini_model.generate_content(new_context)\n",
    "                text = response.text.strip()\n",
    "                # json_response = json.loads(text)\n",
    "                # result = json_response[example_cases]\n",
    "                gemini_responses.append(text)\n",
    "            except:\n",
    "                gemini_responses.append(\"ERROR\")\n",
    "\n",
    "            time.sleep(60)\n",
    "            counter += 1\n",
    "        else:\n",
    "            gemini_responses.append(\"NOT SENT\")\n",
    "            \n",
    "    return gemini_responses\n",
    "\n",
    "\n",
    "def segment_image(image_path, text_prompt):\n",
    "    \"\"\"\n",
    "    Segments an image and uses Gemini to verify if segments match the prompt.\n",
    "\n",
    "    Args:\n",
    "      image_path: Path to the image.\n",
    "      text_prompt: Text prompt describing the object to segment.\n",
    "\n",
    "    Returns:\n",
    "      A list of masks, confidence scores, and Gemini's YES/NO responses.\n",
    "    \"\"\"\n",
    "    # Generate masks using your SAM 2 setup\n",
    "    [generated_masks, _scores] = calculate_sam_masks(image_path)\n",
    "\n",
    "    # Calculate similarity between text embedding and mask embeddings\n",
    "    clip_confidences = calculate_clips_similarity(\n",
    "        image_path, generated_masks, text_prompt\n",
    "    )\n",
    "\n",
    "    # Calculate Gemini responses\n",
    "    gemini_responses = calculate_gemini_response(\n",
    "        image_path, generated_masks, text_prompt\n",
    "    )\n",
    "\n",
    "    return generated_masks, clip_confidences, gemini_responses\n",
    "\n",
    "\n",
    "def plot_results(image_path, masks, confidences, gemini_responses):\n",
    "    \"\"\"\n",
    "    Plots the image with the generated masks and confidences.\n",
    "\n",
    "    Args:\n",
    "      image_path: Path to the image.\n",
    "      masks: A list of masks.\n",
    "      confidences: A list of confidence scores.\n",
    "      gemini_responses: A list of Gemini responses.\n",
    "    \"\"\"\n",
    "    image = Image.open(image_path)\n",
    "    num_masks = len(masks)\n",
    "    num_cols = 4\n",
    "    num_rows = (num_masks + num_cols - 1) // num_cols\n",
    "\n",
    "    plt.figure(figsize=(15, 5 * num_rows))\n",
    "    for i, (mask, confidence, gemini_response) in enumerate(\n",
    "        zip(masks, confidences, gemini_responses)\n",
    "    ):\n",
    "        plt.subplot(num_rows, num_cols, i + 1)\n",
    "        plt.imshow(image)\n",
    "        plt.imshow(mask, alpha=0.5)\n",
    "        plt.title(\n",
    "            f\"Mask {i+1}\\n(Confidence: {confidence:.2f})\\n(Gemini: {gemini_response})\"\n",
    "        )\n",
    "        plt.axis(\"off\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Uploaded file 'cars.jpg' as: https://generativelanguage.googleapis.com/v1beta/files/m6dopuxka32j\n",
      "Uploaded file 'red_car.png' as: https://generativelanguage.googleapis.com/v1beta/files/ver712q6buip\n",
      "Uploaded file 'pink_car_0.png' as: https://generativelanguage.googleapis.com/v1beta/files/xrzafuuwf9ed\n",
      "Uploaded file 'pink_car_1.png' as: https://generativelanguage.googleapis.com/v1beta/files/m58i4acbqlo0\n",
      "Uploaded file 'red_car_wheel.png' as: https://generativelanguage.googleapis.com/v1beta/files/a15tors3sqzt\n",
      "Uploaded file 'road.png' as: https://generativelanguage.googleapis.com/v1beta/files/r0p2nqxulpey\n"
     ]
    }
   ],
   "source": [
    "# Load the models\n",
    "clip_model, preprocess = load_clip_model()\n",
    "sam2, mask_generator = load_sam_model()\n",
    "gemini_model, context, example_cases = load_gemini_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "ename": "OutOfMemoryError",
     "evalue": "CUDA out of memory. Tried to allocate 1.47 GiB. GPU 0 has a total capacity of 5.92 GiB of which 734.00 MiB is free. Including non-PyTorch memory, this process has 5.19 GiB memory in use. Of the allocated memory 3.20 GiB is allocated by PyTorch, and 1.92 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mOutOfMemoryError\u001b[0m                          Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[25], line 4\u001b[0m\n\u001b[1;32m      2\u001b[0m image_path \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m./images/cars.jpg\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m      3\u001b[0m text_prompt \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ma pink car\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m----> 4\u001b[0m masks, confidences, gemini_responses \u001b[38;5;241m=\u001b[39m segment_image(image_path, text_prompt)\n\u001b[1;32m      5\u001b[0m plot_results(image_path, masks, confidences, gemini_responses)\n",
      "Cell \u001b[0;32mIn[24], line 326\u001b[0m, in \u001b[0;36msegment_image\u001b[0;34m(image_path, text_prompt)\u001b[0m\n\u001b[1;32m    315\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    316\u001b[0m \u001b[38;5;124;03mSegments an image and uses Gemini to verify if segments match the prompt.\u001b[39;00m\n\u001b[1;32m    317\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    323\u001b[0m \u001b[38;5;124;03m  A list of masks, confidence scores, and Gemini's YES/NO responses.\u001b[39;00m\n\u001b[1;32m    324\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    325\u001b[0m \u001b[38;5;66;03m# Generate masks using your SAM 2 setup\u001b[39;00m\n\u001b[0;32m--> 326\u001b[0m [generated_masks, _scores] \u001b[38;5;241m=\u001b[39m calculate_sam_masks(image_path)\n\u001b[1;32m    328\u001b[0m \u001b[38;5;66;03m# Calculate similarity between text embedding and mask embeddings\u001b[39;00m\n\u001b[1;32m    329\u001b[0m clip_confidences \u001b[38;5;241m=\u001b[39m calculate_clips_similarity(\n\u001b[1;32m    330\u001b[0m     image_path, generated_masks, text_prompt\n\u001b[1;32m    331\u001b[0m )\n",
      "Cell \u001b[0;32mIn[24], line 235\u001b[0m, in \u001b[0;36mcalculate_sam_masks\u001b[0;34m(image_path)\u001b[0m\n\u001b[1;32m    233\u001b[0m image \u001b[38;5;241m=\u001b[39m Image\u001b[38;5;241m.\u001b[39mopen(image_path)\n\u001b[1;32m    234\u001b[0m image_array \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39marray(image)\n\u001b[0;32m--> 235\u001b[0m masks \u001b[38;5;241m=\u001b[39m mask_generator\u001b[38;5;241m.\u001b[39mgenerate(image_array)\n\u001b[1;32m    237\u001b[0m \u001b[38;5;66;03m# Extract masks and scores (adapt this based on SAM 2 output format)\u001b[39;00m\n\u001b[1;32m    238\u001b[0m generated_masks \u001b[38;5;241m=\u001b[39m [mask[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124msegmentation\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;28;01mfor\u001b[39;00m mask \u001b[38;5;129;01min\u001b[39;00m masks]\n",
      "File \u001b[0;32m~/miniconda3/envs/clip/lib/python3.12/site-packages/torch/utils/_contextlib.py:116\u001b[0m, in \u001b[0;36mcontext_decorator.<locals>.decorate_context\u001b[0;34m(*args, **kwargs)\u001b[0m\n\u001b[1;32m    113\u001b[0m \u001b[38;5;129m@functools\u001b[39m\u001b[38;5;241m.\u001b[39mwraps(func)\n\u001b[1;32m    114\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mdecorate_context\u001b[39m(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs):\n\u001b[1;32m    115\u001b[0m     \u001b[38;5;28;01mwith\u001b[39;00m ctx_factory():\n\u001b[0;32m--> 116\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m func(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n",
      "File \u001b[0;32m~/miniconda3/envs/clip/lib/python3.12/site-packages/sam2/automatic_mask_generator.py:196\u001b[0m, in \u001b[0;36mSAM2AutomaticMaskGenerator.generate\u001b[0;34m(self, image)\u001b[0m\n\u001b[1;32m    171\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    172\u001b[0m \u001b[38;5;124;03mGenerates masks for the given image.\u001b[39;00m\n\u001b[1;32m    173\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    192\u001b[0m \u001b[38;5;124;03m         the mask, given in XYWH format.\u001b[39;00m\n\u001b[1;32m    193\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m    195\u001b[0m \u001b[38;5;66;03m# Generate masks\u001b[39;00m\n\u001b[0;32m--> 196\u001b[0m mask_data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_generate_masks(image)\n\u001b[1;32m    198\u001b[0m \u001b[38;5;66;03m# Encode masks\u001b[39;00m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39moutput_mode \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mcoco_rle\u001b[39m\u001b[38;5;124m\"\u001b[39m:\n",
      "File \u001b[0;32m~/miniconda3/envs/clip/lib/python3.12/site-packages/sam2/automatic_mask_generator.py:233\u001b[0m, in \u001b[0;36mSAM2AutomaticMaskGenerator._generate_masks\u001b[0;34m(self, image)\u001b[0m\n\u001b[1;32m    231\u001b[0m data \u001b[38;5;241m=\u001b[39m MaskData()\n\u001b[1;32m    232\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m crop_box, layer_idx \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(crop_boxes, layer_idxs):\n\u001b[0;32m--> 233\u001b[0m     crop_data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_process_crop(image, crop_box, layer_idx, orig_size)\n\u001b[1;32m    234\u001b[0m     data\u001b[38;5;241m.\u001b[39mcat(crop_data)\n\u001b[1;32m    236\u001b[0m \u001b[38;5;66;03m# Remove duplicate masks between crops\u001b[39;00m\n",
      "File \u001b[0;32m~/miniconda3/envs/clip/lib/python3.12/site-packages/sam2/automatic_mask_generator.py:271\u001b[0m, in \u001b[0;36mSAM2AutomaticMaskGenerator._process_crop\u001b[0;34m(self, image, crop_box, crop_layer_idx, orig_size)\u001b[0m\n\u001b[1;32m    269\u001b[0m data \u001b[38;5;241m=\u001b[39m MaskData()\n\u001b[1;32m    270\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m (points,) \u001b[38;5;129;01min\u001b[39;00m batch_iterator(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpoints_per_batch, points_for_image):\n\u001b[0;32m--> 271\u001b[0m     batch_data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_process_batch(\n\u001b[1;32m    272\u001b[0m         points, cropped_im_size, crop_box, orig_size, normalize\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m    273\u001b[0m     )\n\u001b[1;32m    274\u001b[0m     data\u001b[38;5;241m.\u001b[39mcat(batch_data)\n\u001b[1;32m    275\u001b[0m     \u001b[38;5;28;01mdel\u001b[39;00m batch_data\n",
      "File \u001b[0;32m~/miniconda3/envs/clip/lib/python3.12/site-packages/sam2/automatic_mask_generator.py:334\u001b[0m, in \u001b[0;36mSAM2AutomaticMaskGenerator._process_batch\u001b[0;34m(self, points, im_size, crop_box, orig_size, normalize)\u001b[0m\n\u001b[1;32m    332\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpred_iou_thresh \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0.0\u001b[39m:\n\u001b[1;32m    333\u001b[0m     keep_mask \u001b[38;5;241m=\u001b[39m data[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124miou_preds\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mpred_iou_thresh\n\u001b[0;32m--> 334\u001b[0m     data\u001b[38;5;241m.\u001b[39mfilter(keep_mask)\n\u001b[1;32m    336\u001b[0m \u001b[38;5;66;03m# Calculate and filter by stability score\u001b[39;00m\n\u001b[1;32m    337\u001b[0m data[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mstability_score\u001b[39m\u001b[38;5;124m\"\u001b[39m] \u001b[38;5;241m=\u001b[39m calculate_stability_score(\n\u001b[1;32m    338\u001b[0m     data[\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mmasks\u001b[39m\u001b[38;5;124m\"\u001b[39m], \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mmask_threshold, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstability_score_offset\n\u001b[1;32m    339\u001b[0m )\n",
      "File \u001b[0;32m~/miniconda3/envs/clip/lib/python3.12/site-packages/sam2/utils/amg.py:51\u001b[0m, in \u001b[0;36mMaskData.filter\u001b[0;34m(self, keep)\u001b[0m\n\u001b[1;32m     49\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stats[k] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m     50\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(v, torch\u001b[38;5;241m.\u001b[39mTensor):\n\u001b[0;32m---> 51\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stats[k] \u001b[38;5;241m=\u001b[39m v[torch\u001b[38;5;241m.\u001b[39mas_tensor(keep, device\u001b[38;5;241m=\u001b[39mv\u001b[38;5;241m.\u001b[39mdevice)]\n\u001b[1;32m     52\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(v, np\u001b[38;5;241m.\u001b[39mndarray):\n\u001b[1;32m     53\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_stats[k] \u001b[38;5;241m=\u001b[39m v[keep\u001b[38;5;241m.\u001b[39mdetach()\u001b[38;5;241m.\u001b[39mcpu()\u001b[38;5;241m.\u001b[39mnumpy()]\n",
      "\u001b[0;31mOutOfMemoryError\u001b[0m: CUDA out of memory. Tried to allocate 1.47 GiB. GPU 0 has a total capacity of 5.92 GiB of which 734.00 MiB is free. Including non-PyTorch memory, this process has 5.19 GiB memory in use. Of the allocated memory 3.20 GiB is allocated by PyTorch, and 1.92 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)"
     ]
    }
   ],
   "source": [
    "# Example usage\n",
    "image_path = \"./images/cars.jpg\"\n",
    "text_prompt = \"a pink car\"\n",
    "masks, confidences, gemini_responses = segment_image(image_path, text_prompt)\n",
    "plot_results(image_path, masks, confidences, gemini_responses)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "clip",
   "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.12.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}