YOLO+SAM2.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "authorship_tag": "ABX9TyMA16JhSaanhXLrnWEUfTcg",
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/AbelAbeb/object-detection-using-yolo-and-sam2/blob/main/YOLO%2BSAM2.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4lmbe1yM0GoA"
      },
      "outputs": [],
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/drive')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "%cd /content/drive/MyDrive/AI_Class/PROJECT_PUB\n",
        "!git clone https://github.com/facebookresearch/sam2.git\n",
        "%cd sam2\n",
        "!pip install -e .\n",
        "\n",
        "%cd /content/drive/MyDrive/PROJECT_PUB/sam2/checkpoints\n",
        "!chmod +x ./download_ckpts.sh\n",
        "\n",
        "!./download_ckpts.sh\n",
        "%cd ..\n",
        "%cd /content/drive/MyDrive/PROJECT_PUB/sam2"
      ],
      "metadata": {
        "id": "bCgCX9aD0msb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import torch\n",
        "from sam2.build_sam import build_sam2_video_predictor\n",
        "\n",
        "checkpoint = \"./checkpoints/sam2.1_hiera_tiny.pt\"\n",
        "model_cfg = \"configs/sam2.1/sam2.1_hiera_t.yaml\"\n",
        "predictor = build_sam2_video_predictor(model_cfg, checkpoint)"
      ],
      "metadata": {
        "id": "3O28o0S601TH"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def show_mask(mask, ax, obj_id=None, random_color=False):\n",
        "    if random_color:\n",
        "        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)\n",
        "    else:\n",
        "        cmap = plt.get_cmap(\"tab10\")\n",
        "        cmap_idx = 0 if obj_id is None else obj_id\n",
        "        color = np.array([*cmap(cmap_idx)[:3], 0.6])\n",
        "    h, w = mask.shape[-2:]\n",
        "    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)\n",
        "    ax.imshow(mask_image)\n",
        "\n",
        "\n",
        "def show_points(coords, labels, ax, marker_size=200):\n",
        "    pos_points = coords[labels==1]\n",
        "    neg_points = coords[labels==0]\n",
        "    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
        "    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white', linewidth=1.25)\n",
        "\n",
        "\n",
        "def show_box(box, ax):\n",
        "    x0, y0 = box[0], box[1]\n",
        "    w, h = box[2] - box[0], box[3] - box[1]\n",
        "    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))"
      ],
      "metadata": {
        "id": "IL1l6Cv705Jz"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import numpy as np\n",
        "import torch\n",
        "import matplotlib.pyplot as plt\n",
        "from PIL import Image"
      ],
      "metadata": {
        "id": "CKptaYCy08qM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "torch.autocast(\"cuda\", dtype=torch.bfloat16).__enter__()\n",
        "# turn on tfloat32 for Ampere GPUs (https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices)\n",
        "if torch.cuda.get_device_properties(0).major >= 8:\n",
        "    torch.backends.cuda.matmul.allow_tf32 = True\n",
        "    torch.backends.cudnn.allow_tf32 = True"
      ],
      "metadata": {
        "id": "Qg7hrQDl0_gb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "#%cd /content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/Video\n",
        "from supervision.assets import download_assets, VideoAssets\n",
        "import supervision as sv\n",
        "SOURCE_VIDEO = download_assets(VideoAssets.BASKETBALL)\n",
        "#SOURCE_VIDEO = '/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/Video/video70.mp4'\n",
        "sv.VideoInfo.from_video_path(SOURCE_VIDEO)"
      ],
      "metadata": {
        "id": "oK-fd_vA1EYM"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "0mryCX5I1io-"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Generating the images for the first time from the given videos**"
      ],
      "metadata": {
        "id": "uua6_m6P1lT6"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "SCALE_FACTOR = 0.5\n",
        "START_IDX = 0\n",
        "END_IDX = 476\n",
        "\n",
        "from pathlib import Path\n",
        "#%cd /content/drive/MyDrive/AI_Class/PROJECT_PUB/sam2/custom_dataset\n",
        "#SOURCE_FRAMES = Path(HOME) / Path(SOURCE_VIDEO).stem\n",
        "SOURCE_FRAMES = Path('/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/images1')\n",
        "SOURCE_FRAMES.mkdir(parents=True, exist_ok=True)\n",
        "SOURCE_VIDEO ='/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/Video/video70.mp4'\n",
        "frames_generator = sv.get_video_frames_generator(SOURCE_VIDEO, start_idx = START_IDX, end_idx = END_IDX)\n",
        "images_sink = sv.ImageSink(\n",
        "    target_dir_path=SOURCE_FRAMES.as_posix(),\n",
        "    overwrite=True,\n",
        "    image_name_pattern=\"{:05d}.jpeg\"\n",
        ")\n",
        "\n",
        "with images_sink:\n",
        "    for frame in frames_generator:\n",
        "        frame = sv.scale_image(frame, SCALE_FACTOR)\n",
        "        images_sink.save_image(frame)\n",
        "\n",
        "TARGET_VIDEO = f\"{Path(SOURCE_VIDEO).stem}-result.mp4\"\n",
        "SOURCE_FRAME_PATHS = sorted(sv.list_files_with_extensions(SOURCE_FRAMES.as_posix(), extensions=[\"jpeg\"]))"
      ],
      "metadata": {
        "id": "JPvtn1uC1JX4"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "If the images is already geneated we can use the followind code"
      ],
      "metadata": {
        "id": "vhW7ZmpO1z5m"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "from pathlib import Path\n",
        "SOURCE_FRAMES = Path('/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/images')\n",
        "#SOURCE_FRAMES.mkdir(parents=True, exist_ok=True)\n",
        "SOURCE_VIDEO ='/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/Video/basketball-1.mp4'\n",
        "TARGET_VIDEO = f\"{Path('/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/masked_video').stem}-masked.mp4\"\n",
        "SOURCE_FRAME_PATHS = sorted(sv.list_files_with_extensions(SOURCE_FRAMES.as_posix(), extensions=[\"jpeg\"]))"
      ],
      "metadata": {
        "id": "q-GxCjQV18bX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# The variable VIDEO_FRAMES_DIRECTORY_PATH should be a string\n",
        "VIDEO_FRAMES_DIRECTORY_PATH = '/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/images1'\n",
        "inference_state = predictor.init_state(VIDEO_FRAMES_DIRECTORY_PATH)"
      ],
      "metadata": {
        "id": "8JmBnqU92Arl"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "predictor.reset_state(inference_state)"
      ],
      "metadata": {
        "id": "lMAJda0Z2KoQ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**Now time to prompt on the images. let see easy prompting way using UI**"
      ],
      "metadata": {
        "id": "kPCkFz7A2ROc"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import torch\n",
        "import base64\n",
        "\n",
        "import numpy as np\n",
        "import supervision as sv\n",
        "\n",
        "from pathlib import Path\n",
        "from supervision.assets import download_assets, VideoAssets\n",
        "from sam2.build_sam import build_sam2_video_predictor\n",
        "\n",
        "IS_COLAB = True\n",
        "\n",
        "if IS_COLAB:\n",
        "    from google.colab import output\n",
        "    output.enable_custom_widget_manager()\n",
        "\n",
        "from jupyter_bbox_widget import BBoxWidget"
      ],
      "metadata": {
        "id": "FiAgwH59245D"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [],
      "metadata": {
        "id": "_cML9NXi2-aZ"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def encode_image(filepath):\n",
        "    with open(filepath, 'rb') as f:\n",
        "        image_bytes = f.read()\n",
        "    encoded = str(base64.b64encode(image_bytes), 'utf-8')\n",
        "    return \"data:image/jpg;base64,\"+encoded"
      ],
      "metadata": {
        "id": "sTpTCfci2jq7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**NOTE:** SAM2 allows tracking multiple objects at once. Update the `OBJECTS` list if you want to change the list of tracked objects."
      ],
      "metadata": {
        "id": "HQhnN6Ob2t0r"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "OBJECTS = ['ball', 'player-1', 'player-2']"
      ],
      "metadata": {
        "id": "fLwCoZY72vKi"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "\n",
        "**NOTE:** Let's choose the index of the reference frame that we will use to annotate the objects we are looking for."
      ],
      "metadata": {
        "id": "7v5x8LS62_Ww"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "FRAME_IDX = 100\n",
        "SOURCE_FRAMES = '/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/images'\n",
        "FRAME_PATH = Path(SOURCE_FRAMES) / f\"{FRAME_IDX:05d}.jpeg\"\n",
        "\n",
        "widget = BBoxWidget(classes=OBJECTS)\n",
        "widget.image = encode_image(FRAME_PATH)\n",
        "widget"
      ],
      "metadata": {
        "id": "_cmAGi4_3DZB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "widget.bboxes"
      ],
      "metadata": {
        "id": "fqAkWz2a3Obx"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "**NOTE:** The widget we are using stores annotations in a format that is inconsistent with SAM2's requirements. We parse them and then pass them to SAM2 via the `add_new_points` method. Each of the objects we track must be passed via a separate `add_new_points` call. It is important to specify `frame_idx` each time - the index of the frame to which the annotations relate, and `obj_id` - the ID of the object to which the annotations relate."
      ],
      "metadata": {
        "id": "uFeyazGW3ggS"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "default_box = [\n",
        "    {'x': 705, 'y': 302, 'width': 0, 'height': 0, 'label': 'ball'},\n",
        "    {'x': 587, 'y': 300, 'width': 0, 'height': 0, 'label': 'player-1'},\n",
        "    {'x': 753, 'y': 267, 'width': 0, 'height': 0, 'label': 'player-2'}\n",
        "]\n",
        "\n",
        "boxes = widget.bboxes if widget.bboxes else default_box\n",
        "\n",
        "for object_id, label in enumerate(OBJECTS, start=1):\n",
        "    boxes = [box for box in widget.bboxes if box['label'] == label]\n",
        "\n",
        "    if len(boxes) == 0:\n",
        "        continue\n",
        "\n",
        "    points = np.array([\n",
        "        [\n",
        "            box['x'],\n",
        "            box['y']\n",
        "        ] for box in boxes\n",
        "    ], dtype=np.float32)\n",
        "    labels = np.ones(len(points))\n",
        "\n",
        "    _, object_ids, mask_logits = predictor.add_new_points(\n",
        "        inference_state=inference_state,\n",
        "        frame_idx=FRAME_IDX,\n",
        "        obj_id=object_id,\n",
        "        points=points,\n",
        "        labels=labels,\n",
        "    )"
      ],
      "metadata": {
        "id": "c8xps0I03ePX"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Video inference\n",
        "\n",
        "**NOTE:** To apply our point prompts to all video frames, we use the `propagate_in_video` generator. Each call returns `frame_idx` - the index of the current frame, `object_ids` - IDs of objects detected in the frame, and `mask_logits` - corresponding `object_ids` logit values, which we can convert to masks using thresholding."
      ],
      "metadata": {
        "id": "yQmNvRLX3o5y"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "*Here sam2 will go through the given video and save the masked images in the specified folder *"
      ],
      "metadata": {
        "id": "MBmYv9l430Mp"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "video_info = sv.VideoInfo.from_video_path(SOURCE_VIDEO)\n",
        "video_info.width = int(video_info.width * SCALE_FACTOR)\n",
        "video_info.height = int(video_info.height * SCALE_FACTOR)\n",
        "annotated_frames_dir = '/content/drive/MyDrive/AI_Class/PROJECT_PUB/sam2/custom_dataset/Masked_images'\n",
        "COLORS = ['#FF1493', '#00BFFF', '#FF6347', '#FFD700']\n",
        "mask_annotator = sv.MaskAnnotator(\n",
        "    color=sv.ColorPalette.from_hex(COLORS),\n",
        "    color_lookup=sv.ColorLookup.CLASS)\n",
        "\n",
        "frame_sample = []\n",
        "frame_paths = []\n",
        "with sv.VideoSink(Path(TARGET_VIDEO).as_posix(), video_info=video_info) as sink:\n",
        "    for frame_idx, object_ids, mask_logits in predictor.propagate_in_video(inference_state):\n",
        "        frame_path = SOURCE_FRAME_PATHS[frame_idx]\n",
        "        frame_paths.append(frame_path)\n",
        "        frame = cv2.imread(frame_path)\n",
        "        masks = (mask_logits > 0.0).cpu().numpy()\n",
        "        masks = np.squeeze(masks).astype(bool)\n",
        "\n",
        "        detections = sv.Detections(\n",
        "            xyxy=sv.mask_to_xyxy(masks=masks),\n",
        "            mask=masks,\n",
        "            class_id=np.array(object_ids)\n",
        "        )\n",
        "\n",
        "        annotated_frame = mask_annotator.annotate(scene=frame.copy(), detections=detections)\n",
        "\n",
        "        sink.write_frame(annotated_frame)\n",
        "        #saving specific annotated frames\n",
        "        #if frame_idx % video_info.fps == 0:\n",
        "        #   frame_sample.append(annotated_frame)\n",
        "        #saving all masked frames\n",
        "        frame_sample.append(annotated_frame)"
      ],
      "metadata": {
        "id": "Nw3c4HOy3lJW"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Now time to integrate the YOLO model and sam2 model. here we will get the bounding box of the objects from yolo for each images the sam2 try to sagment the objects based on the given boxes."
      ],
      "metadata": {
        "id": "L6uHvKmM4wkD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import os\n",
        "from ultralytics import YOLO\n",
        "\n",
        "def extract_detection_info(frame_folder, model_path=\"yolov8s.pt\"):\n",
        "    \"\"\"\n",
        "    Extract detection information using YOLO model without visualization\n",
        "\n",
        "    Args:\n",
        "        frame_folder (str): Path to folder containing image frames\n",
        "        model_path (str): Path to YOLO model weights\n",
        "\n",
        "    Returns:\n",
        "        list: List of dictionaries containing detection information\n",
        "    \"\"\"\n",
        "    # Initialize YOLO model\n",
        "    yolo = YOLO(model_path)\n",
        "\n",
        "    all_detections = []  # Store all detections\n",
        "\n",
        "    # Process each frame\n",
        "    for frame_file in os.listdir(frame_folder):\n",
        "        if not frame_file.lower().endswith(('.png', '.jpg', '.jpeg')):\n",
        "            continue\n",
        "\n",
        "        frame_path = os.path.join(frame_folder, frame_file)\n",
        "        frame_id = os.path.splitext(frame_file)[0]\n",
        "\n",
        "        # Load frame\n",
        "        frame = cv2.imread(frame_path)\n",
        "        if frame is None:\n",
        "            print(f\"Could not load frame: {frame_path}\")\n",
        "            continue\n",
        "\n",
        "        # Perform detection\n",
        "        results = yolo(frame)\n",
        "\n",
        "        # Process each detection\n",
        "        for result in results:\n",
        "            boxes = result.boxes  # Boxes object for bbox outputs\n",
        "            for box in boxes:\n",
        "                # Get box coordinates\n",
        "                bbox = box.xyxy[0].cpu().numpy()  # get box coordinates in (x1, y1, x2, y2) format\n",
        "\n",
        "                # Get class information\n",
        "                cls = int(box.cls[0].item())  # class id\n",
        "                conf = float(box.conf[0].item())  # confidence score\n",
        "\n",
        "                # Store detection information\n",
        "                detection_info = {\n",
        "                    \"frame_id\": frame_id,\n",
        "                    \"object_id\": cls,\n",
        "                    \"bbox\": bbox.tolist(),  # convert to list for easier handling\n",
        "                    \"confidence\": conf\n",
        "                }\n",
        "\n",
        "                all_detections.append(detection_info)\n",
        "\n",
        "    return all_detections\n",
        "\n",
        "# Usage example\n",
        "if __name__ == \"__main__\":\n",
        "    frame_folder = \"/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/images\"\n",
        "    detections = extract_detection_info(frame_folder, \"yolov10s.pt\")\n",
        "\n",
        "    # Print all detections\n",
        "    print(f\"\\nTotal number of detections across all frames: {len(detections)}\")\n",
        "    for detection in detections:\n",
        "        print(\"\\nDetection:\")\n",
        "        print(detection)\n",
        "\n",
        "        boxes = [detection['bbox'][0],detection['bbox'][1],detection['bbox'][0] + detection['bbox'][2],detection['bbox'][1]+ detection['bbox'][3]]\n",
        "        #for i in range(len(boxes)):\n",
        "        # Pass input_boxes instead of boxes\n",
        "        _, object_ids, mask_logits = predictor.add_new_points_or_box(\n",
        "            inference_state=inference_state,\n",
        "            frame_idx=int(detection['frame_id']),\n",
        "            #points=points,\n",
        "            #labels=labels,\n",
        "            obj_id=detection['object_id'],\n",
        "            box=detection['bbox'],  # Pass the bounding boxes as input_boxes\n",
        "        )\n",
        "\n"
      ],
      "metadata": {
        "id": "bN0c5OYn5b3P"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "problem of this method is with in a given image if we have the same object id/class id for 2 different but the same class object, sam2 will mask one of them. here we have different bounding box but the same object id/ class id which is a problem for the sam2. because same need unique object id to track and mask the object"
      ],
      "metadata": {
        "id": "FNwWoN-d5elv"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import os\n",
        "from ultralytics import YOLO\n",
        "\n",
        "def extract_detection_info(frame_folder, model_path=\"yolov8s.pt\"):\n",
        "    \"\"\"\n",
        "    Extract detection information using YOLO model without visualization\n",
        "\n",
        "    Args:\n",
        "        frame_folder (str): Path to folder containing image frames\n",
        "        model_path (str): Path to YOLO model weights\n",
        "\n",
        "    Returns:\n",
        "        list: List of dictionaries containing detection information\n",
        "    \"\"\"\n",
        "    # Initialize YOLO model\n",
        "    yolo = YOLO(model_path)\n",
        "\n",
        "    all_detections = []  # Store all detections\n",
        "\n",
        "    # Process each frame\n",
        "    for frame_file in os.listdir(frame_folder):\n",
        "        if not frame_file.lower().endswith(('.png', '.jpg', '.jpeg')):\n",
        "            continue\n",
        "\n",
        "        frame_path = os.path.join(frame_folder, frame_file)\n",
        "        frame_id = os.path.splitext(frame_file)[0]\n",
        "\n",
        "        # Load frame\n",
        "        frame = cv2.imread(frame_path)\n",
        "        if frame is None:\n",
        "            print(f\"Could not load frame: {frame_path}\")\n",
        "            continue\n",
        "\n",
        "        # Perform detection\n",
        "        results = yolo(frame)\n",
        "\n",
        "        # Process each detection\n",
        "        for result in results:\n",
        "            boxes = result.boxes  # Boxes object for bbox outputs\n",
        "            for box in boxes:\n",
        "                # Get box coordinates\n",
        "                bbox = box.xyxy[0].cpu().numpy()  # get box coordinates in (x1, y1, x2, y2) format\n",
        "\n",
        "                # Get class information\n",
        "                cls = int(box.cls[0].item())  # class id\n",
        "                conf = float(box.conf[0].item())  # confidence score\n",
        "\n",
        "                # Store detection information\n",
        "                detection_info = {\n",
        "                    \"frame_id\": frame_id,\n",
        "                    \"object_id\": cls,\n",
        "                    \"bbox\": bbox.tolist(),  # convert to list for easier handling\n",
        "                    \"confidence\": conf\n",
        "                }\n",
        "\n",
        "                all_detections.append(detection_info)\n",
        "\n",
        "    return all_detections\n",
        "\n",
        "# Usage example\n",
        "if __name__ == \"__main__\":\n",
        "    frame_folder = \"/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/images\"\n",
        "    detections = extract_detection_info(frame_folder, \"yolov10s.pt\")\n",
        "\n",
        "    # Print all detections\n",
        "    print(f\"\\nTotal number of detections across all frames: {len(detections)}\")\n",
        "    for detection in detections:\n",
        "        print(\"\\nDetection:\")\n",
        "        print(detection)\n",
        "\n",
        "        boxes = [detection['bbox'][0],detection['bbox'][1],detection['bbox'][0] + detection['bbox'][2],detection['bbox'][1]+ detection['bbox'][3]]\n",
        "        #for i in range(len(boxes)):\n",
        "        # Pass input_boxes instead of boxes\n",
        "        _, object_ids, mask_logits = predictor.add_new_points_or_box(\n",
        "            inference_state=inference_state,\n",
        "            frame_idx=int(detection['frame_id']),\n",
        "            #points=points,\n",
        "            #labels=labels,\n",
        "            obj_id=detection['object_id'],\n",
        "            box=detection['bbox'],  # Pass the bounding boxes as input_boxes\n",
        "        )\n",
        "\n"
      ],
      "metadata": {
        "id": "WMjayiB85fTT"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "SCALE_FACTOR = 0.5\n",
        "SOURCE_VIDEO = Path('/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/Video/basketball-1.mp4')\n",
        "TARGET_VIDEO = Path(f\"{Path(SOURCE_VIDEO).stem}-boxInput.mp4\")\n",
        "SOURCE_FRAMES = Path('/content/drive/MyDrive/PROJECT_PUB/sam2/custom_dataset/images')\n",
        "SOURCE_FRAME_PATHS = sorted(sv.list_files_with_extensions(SOURCE_FRAMES.as_posix(), extensions=[\"jpeg\"]))\n",
        "\n",
        "import cv2\n",
        "video_info = sv.VideoInfo.from_video_path(SOURCE_VIDEO)\n",
        "video_info.width = int(video_info.width * SCALE_FACTOR)\n",
        "video_info.height = int(video_info.height * SCALE_FACTOR)\n",
        "\n",
        "COLORS = ['#FF1493', '#00BFFF', '#FF6347', '#FFD700']\n",
        "mask_annotator = sv.MaskAnnotator(\n",
        "    color=sv.ColorPalette.from_hex(COLORS),\n",
        "    color_lookup=sv.ColorLookup.CLASS)\n",
        "\n",
        "frame_sample = []\n",
        "\n",
        "with sv.VideoSink(TARGET_VIDEO.as_posix(), video_info=video_info) as sink:\n",
        "    for frame_idx, object_ids, mask_logits in predictor.propagate_in_video(inference_state):\n",
        "        frame_path = SOURCE_FRAME_PATHS[frame_idx]\n",
        "        frame = cv2.imread(frame_path)\n",
        "        masks = (mask_logits > 0.0).cpu().numpy()\n",
        "        masks = np.squeeze(masks).astype(bool)\n",
        "\n",
        "        detections = sv.Detections(\n",
        "            xyxy=sv.mask_to_xyxy(masks=masks),\n",
        "            mask=masks,\n",
        "            class_id=np.array(object_ids)\n",
        "        )\n",
        "\n",
        "        annotated_frame = mask_annotator.annotate(scene=frame.copy(), detections=detections)\n",
        "\n",
        "        sink.write_frame(annotated_frame)\n",
        "        if frame_idx % video_info.fps == 0:\n",
        "            frame_sample.append(annotated_frame)"
      ],
      "metadata": {
        "id": "SL3cONOD6ZJN"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "now we will going to solve the problem. what we have to do here is, we have to give unique id if 2 or more objects of the same class appeared in the image before we feed to the sam points to be tracked in the video."
      ],
      "metadata": {
        "id": "8FIqmk_i6wGb"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import cv2\n",
        "import os\n",
        "from ultralytics import YOLO\n",
        "from collections import defaultdict\n",
        "\n",
        "def extract_detection_info(frame_folder, model_path=\"yolov10s.pt\"):\n",
        "    \"\"\"\n",
        "    Extract detection information using YOLO model with tracking\n",
        "\n",
        "    Args:\n",
        "        frame_folder (str): Path to folder containing image frames\n",
        "        model_path (str): Path to YOLO model weights\n",
        "\n",
        "    Returns:\n",
        "        list: List of dictionaries containing detection information with track IDs\n",
        "    \"\"\"\n",
        "    # Initialize YOLO model with tracking\n",
        "    yolo = YOLO(model_path)\n",
        "\n",
        "    all_detections = []  # Store all detections\n",
        "\n",
        "    # Process each frame\n",
        "    for frame_file in sorted(os.listdir(frame_folder)):  # Sort to ensure consistent frame order\n",
        "        if not frame_file.lower().endswith(('.png', '.jpg', '.jpeg')):\n",
        "            continue\n",
        "\n",
        "        frame_path = os.path.join(frame_folder, frame_file)\n",
        "        frame_id = os.path.splitext(frame_file)[0]\n",
        "\n",
        "        # Load frame\n",
        "        frame = cv2.imread(frame_path)\n",
        "        if frame is None:\n",
        "            print(f\"Could not load frame: {frame_path}\")\n",
        "            continue\n",
        "\n",
        "        # Perform detection with tracking enabled\n",
        "        results = yolo.track(frame, persist=True, tracker=\"botsort.yaml\")  # Enable tracking\n",
        "\n",
        "        if results and len(results) > 0:\n",
        "            # Process each detection\n",
        "            for result in results:\n",
        "                if not hasattr(result, 'boxes'):\n",
        "                    continue\n",
        "\n",
        "                boxes = result.boxes\n",
        "                for box in boxes:\n",
        "                    # Get box coordinates\n",
        "                    bbox = box.xyxy[0].cpu().numpy()\n",
        "\n",
        "                    # Get class information\n",
        "                    cls = int(box.cls[0].item())\n",
        "                    conf = float(box.conf[0].item())\n",
        "\n",
        "                    # Get track ID (if available)\n",
        "                    track_id = None\n",
        "                    if hasattr(box, 'id'):\n",
        "                        track_id = int(box.id.item())\n",
        "\n",
        "                    # Store detection information\n",
        "                    detection_info = {\n",
        "                        \"frame_id\": frame_id,\n",
        "                        \"class_id\": cls,\n",
        "                        \"track_id\": track_id,  # This will be unique for each instance\n",
        "                        \"bbox\": bbox.tolist(),\n",
        "                        \"confidence\": conf\n",
        "                    }\n",
        "\n",
        "                    all_detections.append(detection_info)\n",
        "\n",
        "    return all_detections\n",
        "\n",
        "# Usage example\n",
        "if __name__ == \"__main__\":\n",
        "    frame_folder = \"/content/drive/MyDrive/AI_Class/PROJECT_PUB/sam2/custom_dataset/images\"\n",
        "    detections = extract_detection_info(frame_folder, \"yolov10s.pt\")\n",
        "\n",
        "    # Print all detections\n",
        "    print(f\"\\nTotal number of detections across all frames: {len(detections)}\")\n",
        "\n",
        "    # Group detections by frame to see multiple instances in each frame\n",
        "    frame_detections = defaultdict(list)\n",
        "    for detection in detections:\n",
        "        frame_detections[detection['frame_id']].append(detection)\n",
        "\n",
        "    # Print detections grouped by frame\n",
        "    for frame_id, frame_dets in frame_detections.items():\n",
        "        print(f\"\\nFrame {frame_id}:\")\n",
        "        for det in frame_dets:\n",
        "            print(f\"  Class {det['class_id']}, Track ID {det['track_id']}, Confidence {det['confidence']:.2f}\")\n",
        "\n",
        "        # Use with predictor\n",
        "        for detection in frame_dets:\n",
        "            boxes = detection['bbox']\n",
        "\n",
        "            # Use track_id as the object identifier if available, otherwise use class_id\n",
        "            obj_id = detection['track_id'] if detection['track_id'] is not None else detection['class_id']\n",
        "\n",
        "            _, object_ids, mask_logits = predictor.add_new_points_or_box(\n",
        "                inference_state=inference_state,\n",
        "                frame_idx=int(detection['frame_id']),\n",
        "                obj_id=obj_id,\n",
        "                box=boxes,\n",
        "            )"
      ],
      "metadata": {
        "id": "dJgyTb2h7O8S"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "video_info = sv.VideoInfo.from_video_path(SOURCE_VIDEO)\n",
        "video_info.width = int(video_info.width * SCALE_FACTOR)\n",
        "video_info.height = int(video_info.height * SCALE_FACTOR)\n",
        "annotated_frames_dir = '/content/drive/MyDrive/AI_Class/PROJECT_PUB/sam2/custom_dataset/Masked_images'\n",
        "COLORS = ['#FF1493', '#00BFFF', '#FF6347', '#FFD700']\n",
        "mask_annotator = sv.MaskAnnotator(\n",
        "    color=sv.ColorPalette.from_hex(COLORS),\n",
        "    color_lookup=sv.ColorLookup.CLASS)\n",
        "\n",
        "frame_sample = []\n",
        "frame_paths = []\n",
        "# Define the class mapping\n",
        "# class_mapping = {\n",
        "#     1: 'ball',\n",
        "#     2: 'player-1',\n",
        "#     3: 'player-2'\n",
        "#     # Add more mappings if necessary\n",
        "#}\n",
        "with sv.VideoSink(Path(TARGET_VIDEO).as_posix(), video_info=video_info) as sink:\n",
        "    for frame_idx, object_ids, mask_logits in predictor.propagate_in_video(inference_state):\n",
        "        frame_path = SOURCE_FRAME_PATHS[frame_idx]\n",
        "        frame_paths.append(frame_path)\n",
        "        frame = cv2.imread(frame_path)\n",
        "        masks = (mask_logits > 0.0).cpu().numpy()\n",
        "        masks = np.squeeze(masks).astype(bool)\n",
        "\n",
        "        detections = sv.Detections(\n",
        "            xyxy=sv.mask_to_xyxy(masks=masks),\n",
        "            mask=masks,\n",
        "            class_id=np.array(object_ids)\n",
        "        )\n",
        "\n",
        "        #Get class names for each detection\n",
        "        #class_names = [class_mapping[id] for id, mask in zip(detections.class_id, detections.mask) if np.any(mask)]\n",
        "\n",
        "        # Print or store class names as needed\n",
        "        #print(f\"Frame {frame_idx}: {class_names}\")\n",
        "\n",
        "        annotated_frame = mask_annotator.annotate(scene=frame.copy(), detections=detections)\n",
        "        # Save the annotated frame to the specified directory\n",
        "        frame_path = os.path.join(annotated_frames_dir, f\"{frame_idx:05d}.jpg\")  # Example filename\n",
        "        cv2.imwrite(frame_path, annotated_frame)\n",
        "        frame_paths.append(frame_path)\n",
        "\n",
        "         # Write the annotated frame to the video sink\n",
        "\n",
        "        sink.write_frame(annotated_frame)\n",
        "        #saving specific annotated frames\n",
        "        #if frame_idx % video_info.fps == 0:\n",
        "        #   frame_sample.append(annotated_frame)\n",
        "        #saving all masked frames\n",
        "        frame_sample.append(annotated_frame)\n"
      ],
      "metadata": {
        "id": "ExACv-2N78X0"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Visualize the result"
      ],
      "metadata": {
        "id": "8bmUNrZr8ChY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sv.plot_images_grid(\n",
        "    images=frame_sample[:20],\n",
        "    grid_size=(5, 5)\n",
        ")"
      ],
      "metadata": {
        "id": "cxGoAjwP8A23"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}