{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Use YOLO model from Ultralytics\n",
"\n",
"This notebook demonstrates how to use a task agent to pre-label images with predictions.\n",
"Here we'll use a bounding box prediction model.\n",
"\n",
"Before we start, let's get installations and authentication out of the way.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 1: Set up environment"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Installation\n",
"\n",
"Please ensure that you have the `encord-agents` library installed:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!python -m pip install encord-agents"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Authentication\n",
"\n",
"The library authenticates via ssh-keys. Below, is a code cell for setting the `ENCORD_SSH_KEY` environment variable. It should contain the raw content of your private ssh key file.\n",
"\n",
"If you have not yet setup an ssh key, please follow the [documentation](https://agents-docs.encord.com/authentication/).\n",
"\n",
"> 💡 **Colab users**: In colab, you can set the key once in the secrets in the left sidebar and load it in new notebooks with\n",
"> ```python\n",
"> from google.colab import userdata\n",
"> key_content = userdata.get(\"ENCORD_SSH_KEY\")\n",
"> ```"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from google.colab import userdata\n",
"\n",
"key_contet = userdata.get(\"ENCORD_SSH_KEY\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"\n",
"os.environ[\"ENCORD_SSH_KEY\"] = key_contet\n",
"# or you can set a path to a file\n",
"# os.environ[\"ENCORD_SSH_KEY_FILE\"] = \"/path/to/your/private/key\""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### [Alternative] Temporary Key\n",
"There's also the option of generating a temporary (fresh) ssh key pair via the code cell below.\n",
"Please follow the instructions printed when executing the code."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# ⚠️ Safe to skip if you have authenticated already\n",
"import os\n",
"\n",
"from encord_agents.utils.colab import generate_public_private_key_pair_with_instructions\n",
"\n",
"private_key_path, public_key_path = generate_public_private_key_pair_with_instructions()\n",
"os.environ[\"ENCORD_SSH_KEY_FILE\"] = private_key_path.as_posix()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 2: Set up your Ontology\n",
"\n",
"Create an Ontology that matches the expected output of your pre-labeling agent.\n",
"For example, if your model predicts classes `surfboard`, `person`, and `car`, then the ontology should look like this: Our DETR model predicts more objects but we'll focus on the car predictions in this example"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
" \n",
" Figure 1: Project ontology.\n",
""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### 2.1 Define an Ontology Map\n",
"We need to translate the model predictions so that they are paired against the respective Encord ontology item. This is easiest done via the featureNodeHash of the target. This can be found in the app either via the Ontology preview JSON or via using the SDK.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"ontology_map = {\"car\": \"uFmVW/cr\"}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 3: Define a model for predictions\n",
"\n",
"We will define a model which predicts labels, bounding boxes, and confidences.\n",
"We'll use the model to predict objects on images\n",
"\n",
"We'll use the YOLOv11 model from Ultralytics following:\n",
"https://docs.ultralytics.com/modes/predict/"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip install ultralytics"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from ultralytics import YOLO\n",
"\n",
"model = YOLO(\"yolo11n.pt\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": []
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"[📖 here](https://docs.encord.com/platform-documentation/GettingStarted/gettingstarted-create-ontology) is the documentation for creating ontologies."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from dataclasses import dataclass\n",
"\n",
"import numpy as np\n",
"from encord.objects.coordinates import BoundingBoxCoordinates\n",
"from numpy.typing import NDArray\n",
"\n",
"\n",
"# Data class to hold predictions from our model\n",
"@dataclass\n",
"class ModelPrediction:\n",
" featureHash: str\n",
" coords: BoundingBoxCoordinates\n",
" conf: float\n",
"\n",
"\n",
"def YOLO_predict(image: NDArray[np.uint8]) -> list[ModelPrediction]:\n",
" print(image, image.shape, type(image))\n",
" outputs = model(image)\n",
"\n",
" print(outputs)\n",
" for output in outputs:\n",
" clses = [ontology_map[x] for x in output.boxes.cls.tolist()]\n",
" print(clses)\n",
" confs = output.boxes.conf.tolist()\n",
" boxes = output.boxes.xywhn.tolist()\n",
" encord_boxes = [\n",
" BoundingBoxCoordinates(height=x[3], width=x[2], top_left_x=x[0], top_left_y=x[1]) for x in boxes\n",
" ]\n",
" return [ModelPrediction(featureHash=x, coords=y, conf=z) for x, y, z in zip(clses, encord_boxes, confs)]\n",
"\n",
"\n",
"agent = YOLO_predict"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 4: Create a Workflow with a pre-labeling agent node\n",
"\n",
"Create a project in the Encord platform that has a Workflow that includes a pre-labeling agent node before the annotation stage to automatically pre-label tasks with model predictions.\n",
"This node is where we'll hook in our custom code to pre-label the data."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"\n",
" \n",
" Figure 2: Project workflow.\n",
""
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Notice how the workflow has a purple Agent node called \"pre-label.\"\n",
"This node will allow our custom code to run inference over the data before passing it on to the annotation stage.\n",
"\n",
"[📖 here](https://docs.encord.com/platform-documentation/Annotate/annotate-projects/annotate-workflows-and-templates#creating-workflows) is the documentation for creating a workflow with Encord."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Step 5: Define the pre-labelling agent\n",
"\n",
"The following code provides a template for defining an agent that does pre-labeling.\n",
"We assume that the project only contains videos and the we want to do pre-labeling on all frames in each video.\n",
"\n",
"If your agent node is named \"pre-label\" and the pathway to the annotation stage is named \"annotate,\" you will only have to change the `` to your actual project hash to make it work.\n",
"Is your naming, on the other hand, different, then you can update the `stage` parameter of the decorator and the returned string, respectively, to comply with your own setup.\n",
"\n",
"Note that this code uses the [`dep_video_iterator` dependency](../../reference/task_agents.md#encord_agents.tasks.dependencies.dep_video_iterator) to automatically load an iterator of frames as RGB numpy arrays from the video."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from typing import Iterable\n",
"\n",
"import numpy as np\n",
"from encord.objects.ontology_labels_impl import LabelRowV2\n",
"from encord.project import Project\n",
"from numpy.typing import NDArray\n",
"from typing_extensions import Annotated\n",
"\n",
"from encord_agents.core.data_model import Frame\n",
"from encord_agents.tasks import Depends, Runner\n",
"from encord_agents.tasks.dependencies import dep_single_frame\n",
"\n",
"# a. Define a runner that will execute the agent on every task in the agent stage\n",
"runner = Runner(project_hash=\"526c15b2-0dda-456a-952d-b9789c64e79b\")\n",
"\n",
"\n",
"# b. Specify the logic that goes into the \"pre-label\" agent node.\n",
"@runner.stage(stage=\"pre-label\")\n",
"def pre_segment(\n",
" lr: LabelRowV2,\n",
" project: Project,\n",
" frame: Annotated[NDArray[np.uint8], Depends(dep_single_frame)],\n",
") -> str:\n",
" ontology = project.ontology_structure\n",
"\n",
" # d. Predict - we could do batching here to speed up the process\n",
" outputs = agent(frame)\n",
"\n",
" # e. Store the results\n",
" for output in outputs:\n",
" ins = ontology.get_child_by_hash(output.featureHash).create_instance()\n",
" ins.set_for_frames(frames=0, coordinates=output.coords, confidence=output.conf)\n",
"\n",
" lr.add_object_instance(ins)\n",
"\n",
" lr.save()\n",
" return \"annotate\" # Tell where the task should go"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Running the agent\n",
"Now that we've defined the project, workflow, and the agent, it's time to try it out.\n",
"The `runner` object is callable which means that you can just call it to prioritize your tasks."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Run the agent\n",
"runner()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Your agent now assigns labels to the videos and routes them appropriately through the Workflow to the annotation stage.\n",
"As a result, every annotation task should already have pre-existing labels (predictions) included.\n",
"\n",
"> 💡*Hint:* If you were to execute this as a python script, you can run it as a command line interface by putting the above code in an `agents.py` file and replacing\n",
"> ```python\n",
"> runner()\n",
"> ```\n",
"> with\n",
"> ```python\n",
"> if __name__ == \"__main__\":\n",
"> runner.run()\n",
"> ```\n",
"> Which will allow you set, e.g., the project hash via the command line:\n",
"> ```bash\n",
"> python agent.py --project-hash \"...\"\n",
"> ```\n"
]
}
],
"metadata": {
"colab": {
"provenance": []
},
"kernelspec": {
"display_name": "encord-agents-Cw_LL1Rx-py3.11",
"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"
}
},
"nbformat": 4,
"nbformat_minor": 0
}