/autodistill-yolo-world

YOLO World base module for use with Autodistill.

Primary LanguagePython

Autodistill YOLO-World Module

This repository contains the code supporting the YOLO-World base model for use with Autodistill.

YOLO-World, developed by Tencent's AI Lab, is a zero-shot object detection model that can run in close-to-real-time on powerful GPUs. You can provide arbitrary text prompts to detect objects in images, and the model will return bounding boxes and class labels for the objects it finds.

You can use YOLO-World in Autodistill to detect objects.

Read the full Autodistill documentation.

Read the YOLO-World Autodistill documentation.

Installation

To use YOLO-World with autodistill, you need to install the following dependency:

pip3 install autodistill-yolo-world

Quickstart

from autodistill_yolo_world import YOLOWorldModel
from autodistill.detection import CaptionOntology
from autodistill.utils.plot import plot

# define an ontology to map class names to our YOLO-World prompt
# the ontology dictionary has the format {caption: class}
# where caption is the prompt sent to the base model, and class is the label that will
# be saved for that caption in the generated annotations
# then, load the model
base_model = YOLOWorldModel(
    ontology=CaptionOntology(
        {
            "person": "person",
            "a forklift": "forklift"
        }
    )
)

# predict on an image
result = base_model.predict("image.jpeg", confidence=0.1)

plot(
    image=cv2.imread("./image.jpeg"),
    classes=base_model.ontology.classes(),
    detections=result
)

# label a folder of images
base_model.label("./context_images", extension=".jpeg")

License

The YOLO-World model is released under a GPT-3.0 license.

🏆 Contributing

We love your input! Please see the core Autodistill contributing guide to get started. Thank you 🙏 to all our contributors!