images/roboflow_cv_utils_fill_img.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "## Draw Bounding Boxes on Images with the Roboflow Python SDK and Supervision\n",
        "<p>\n",
        "  <a align=\"center\" href=\"\" target=\"_blank\">\n",
        "    <img\n",
        "      width=\"60%\"\n",
        "      src=\"https://raw.githubusercontent.com/roboflow/roboflow-computer-vision-utilities/main/figures/roboflow-cv-utilities-header.png\"\n",
        "    >\n",
        "  </a>\n",
        "</p>"
      ],
      "metadata": {
        "id": "cfsYWzUw55Iy"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Installing Dependencies"
      ],
      "metadata": {
        "id": "pjMiL0i0lZLF"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "!pip install -q roboflow>=1.0.5\n",
        "!pip install -q supervision>=0.6.0"
      ],
      "metadata": {
        "id": "XTCBgNLMlYhI"
      },
      "execution_count": 1,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Setting up our utility functions"
      ],
      "metadata": {
        "id": "ej5qGDcnlZox"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "import os\n",
        "import json\n",
        "import cv2\n",
        "import numpy as np\n",
        "from roboflow import Roboflow\n",
        "\n",
        "import supervision as sv\n",
        "from supervision.draw.color import Color\n",
        "from supervision.draw.color import ColorPalette\n",
        "from supervision import Detections, BoxAnnotator\n",
        "\n",
        "\n",
        "def load_roboflow_model(api_key, workspace_id, project_id, version_number):\n",
        "\n",
        "    # authenticate to your Roboflow account and load your model\n",
        "    rf = Roboflow(api_key=api_key)\n",
        "    project = rf.workspace(workspace_id).project(project_id)\n",
        "    version = project.version(version_number)\n",
        "    model = version.model\n",
        "    \n",
        "    return project, model"
      ],
      "metadata": {
        "id": "9_e5mcsTlTuY"
      },
      "execution_count": 8,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Receiving Model Predictions"
      ],
      "metadata": {
        "id": "d2Pj0IJLlZ_J"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def make_prediction(project, model, image_path, confidence, overlap):\n",
        "\n",
        "    # load the image and make predictions with your model\n",
        "    img = cv2.imread(image_path)\n",
        "    predictions = model.predict(image_path, confidence=confidence, overlap=overlap)\n",
        "    predictions_json = predictions.json()\n",
        "    roboflow_xyxy = np.empty((0, 4))\n",
        "    predicted_classes = []\n",
        "    for bounding_box in predictions:\n",
        "        x1 = bounding_box['x'] - bounding_box['width'] / 2\n",
        "        x2 = bounding_box['x'] + bounding_box['width'] / 2\n",
        "        y1 = bounding_box['y'] - bounding_box['height'] / 2\n",
        "        y2 = bounding_box['y'] + bounding_box['height'] / 2\n",
        "        np.vstack((roboflow_xyxy, [x1, y1, x2, y2]))\n",
        "        predicted_classes.append(bounding_box['class'])\n",
        "        \n",
        "        # class_name = bounding_box['class']\n",
        "        # confidence = bounding_box['confidence']\n",
        "        sv_xyxy = Detections(roboflow_xyxy).from_roboflow(\n",
        "            predictions_json,class_list=list((project.classes).keys()))\n",
        "\n",
        "    return img, predictions_json, sv_xyxy, predicted_classes"
      ],
      "metadata": {
        "id": "OYWgcM3SlElH"
      },
      "execution_count": 9,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### [Supervision](https://github.com/roboflow/supervision/)\n",
        "A set of easy-to-use utils that will come in handy in any computer vision project. Supervision is still in pre-release stage. 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid and may change.\n",
        "<p>\n",
        "  <a align=\"center\" href=\"\" target=\"_blank\">\n",
        "    <img\n",
        "      width=\"60%\"\n",
        "      src=\"https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529\"\n",
        "    >\n",
        "  </a>\n",
        "</p>\n",
        "\n",
        "#### [Roboflow Response Object Format](https://docs.roboflow.com/inference/hosted-api#response-object-format) (Object Detection)\n",
        "Draw Bounding Boxes on Images with the Roboflow Python SDK and Supervision\n",
        "\n",
        "<p>\n",
        "  <a align=\"center\" href=\"\" target=\"_blank\">\n",
        "    <img\n",
        "      width=\"50%\"\n",
        "      src=\"https://raw.githubusercontent.com/roboflow/roboflow-computer-vision-utilities/main/figures/ResponseObjectFormat_JSON.png\"\n",
        "    >\n",
        "  </a>\n",
        "</p>"
      ],
      "metadata": {
        "id": "LPCP26FplcPs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def fill_bounding_boxes(image, sv_xyxy, class_ids, add_labels):\n",
        "\n",
        "    #set add_labels to True to show the label for each object\n",
        "    image_with_boxes = BoxAnnotator(\n",
        "        color=Color.black(), thickness=-1).annotate(image, sv_xyxy, labels=class_ids, skip_label=add_labels)\n",
        "\n",
        "    return image_with_boxes"
      ],
      "metadata": {
        "id": "-sHP4xnklI80"
      },
      "execution_count": 21,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Saving the Image with Bounding Boxes"
      ],
      "metadata": {
        "id": "7WWr6O35lbAj"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def save_image(image, original_image_path, output_directory=\"results\"):\n",
        "\n",
        "    os.makedirs(output_directory, exist_ok=True)\n",
        "    filename = os.path.basename(original_image_path)\n",
        "    output_path = os.path.join(output_directory, f\"result_{filename}\")\n",
        "    cv2.imwrite(output_path, image)\n",
        "    return output_path"
      ],
      "metadata": {
        "id": "2Fap9bCBlMEx"
      },
      "execution_count": 11,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "#### Don't forgot to add your image file(s)!"
      ],
      "metadata": {
        "id": "E3m00687mrP-"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Running the whole thing!"
      ],
      "metadata": {
        "id": "Ac53DIQEli_w"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "lg98goeWlCrQ",
        "outputId": "c41a08e4-9ed8-45d3-dd46-3e98927c5ee7"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "loading Roboflow workspace...\n",
            "loading Roboflow project...\n",
            "{'predictions': [{'x': 1113, 'y': 363, 'width': 334, 'height': 280, 'confidence': 0.9350578784942627, 'class': 'machinery', 'image_path': '/content/construction-2-_mp4-162_jpg.rf.edc08e7528429b4315f59637aaf65cf0.jpg', 'prediction_type': 'ObjectDetectionModel'}, {'x': 307, 'y': 314, 'width': 582, 'height': 460, 'confidence': 0.8053478002548218, 'class': 'machinery', 'image_path': '/content/construction-2-_mp4-162_jpg.rf.edc08e7528429b4315f59637aaf65cf0.jpg', 'prediction_type': 'ObjectDetectionModel'}, {'x': 905, 'y': 272, 'width': 722, 'height': 481, 'confidence': 0.7803330421447754, 'class': 'machinery', 'image_path': '/content/construction-2-_mp4-162_jpg.rf.edc08e7528429b4315f59637aaf65cf0.jpg', 'prediction_type': 'ObjectDetectionModel'}], 'image': {'width': '1280', 'height': '720'}}\n",
            "The image has been processed and saved to results/result_construction-2-_mp4-162_jpg.rf.edc08e7528429b4315f59637aaf65cf0.jpg\n"
          ]
        }
      ],
      "source": [
        "def main():\n",
        "\n",
        "    ## Authentication info to load the model\n",
        "    ## Sample project: https://universe.roboflow.com/roboflow-universe-projects/construction-site-safety/model/25\n",
        "    ## Workspace ID: \"roboflow-universe-projects\", Project ID: \"construction-site-safety\", Version Number: 25\n",
        "    api_key = \"ROBOFLOW_PRIVATE_API_KEY\"\n",
        "    workspace_id = \"ROBOFLOW_WORKSPACE_ID\"\n",
        "    project_id = \"ROBOFLOW_POJECT_ID\"\n",
        "    version_number = \"ROBOFLOW_VERSION_NUMBER\"\n",
        "    project, model = load_roboflow_model(api_key, workspace_id, project_id, version_number)\n",
        "\n",
        "    # Make a prediction on the specified image file\n",
        "    image_path = \"/content/test_images/construction-2-_mp4-13_jpg.rf.dcfc8cd7a7c9f6109ff18cab2b50f6bc.jpg\"\n",
        "    confidence = 40\n",
        "    overlap = 30\n",
        "    image, predictions_json, pred_sv_xyxy, predicted_classes = make_prediction(\n",
        "        project, model, image_path, confidence, overlap)\n",
        "\n",
        "    print(predictions_json)\n",
        "\n",
        "    ## Set add_labels to False to draw class labels on the bounding boxes\n",
        "    add_labels = True\n",
        "    for i in range(len(pred_sv_xyxy)):\n",
        "        image_with_boxes = fill_bounding_boxes(image, pred_sv_xyxy, predicted_classes, add_labels)\n",
        "\n",
        "    # Save the image with bounding boxes for the detected objects drawn on them\n",
        "    output_path = save_image(image_with_boxes, image_path)\n",
        "\n",
        "    print(f\"The image has been processed and saved to {output_path}\")\n",
        "\n",
        "if __name__ == \"__main__\":\n",
        "    main()\n"
      ]
    }
  ]
}