机器学习AI算法工程 | YOLO-World开集目标检测，可直接上手

本文来源公众号“机器学习AI算法工程”，仅用于学术分享，侵权删，干货满满。

原文链接：YOLO-World开集目标检测，可直接上手

1. 前言

        关于Demo：尽管YOLO-World官方给出了在线试用的Demo：https://huggingface.co/spaces/stevengrove/YOLO-World ，但还是不如在自己电脑上离线运行一遍来得直接。恰好，近期Ultralytics（YOLOv8）也新增了对YOLO-World的支持，我们不如直接体验一把。

        关于YOLO-World：

        其是一个使用开放词汇进行目标检测的新框架，且是以YOLOv8框架为detector，所以其特点就继承了YOLO系列，也即轻量、快速、性能好。另外，既然是文本和图片一起作为输入，那么就需要有一个文本embedding的模块，这里用的是CLIP，其将用户输入的词汇列表转换为特征向量，与输入图像一起进行推理。

2. 安装（更新）

        由于YOLO-World主要包括两部分内容：YOLO-style的检测器、用于文本embedding的CLIP。因此，我们需要安装一个Ultralytics库、一个CLIP库。

# 新安装：pip install ultralytics # 已有Ultralytics，更新：pip install -U ultralytics
pip install openai-clip

3. 上手体验

模型下载：

这里选用

yolov8s-world.pt：https://github.com/ultralytics/assets/releases/download/v8.1.0/yolov8s-world.pt

当然，也可以使用其他模型，可自行从Ultralytics的官网下载：。

不设定词汇表，进行推理：

这里没有用到CLIP，因为没有设定词汇，也就不需要进行text embedding了

yolo predict model=yolov8s-world.pt source=path/to/image.jpg imgsz=640

设定词汇，进行推理：

from ultralytics import YOLOif __name__ == '__main__':     # Initialize a YOLO-World model    model = YOLO('yolov8/yolov8s-world.pt')  # or choose yolov8m/l-world.pt     # Define custom classes    model.set_classes(["person"])     # Execute prediction for specified categories on an image    results = model.predict('image_01.jpg')     # Show results    results[0].show()

这里，我们设定了只包括一个词汇的词汇表：["person"]，那么自然，我们也可以根据需要设定多个词汇。由于设定了词汇表，所以就会用到CLIP，它会将各个词汇进行嵌入，转换为text feature。

上述代码中用到了set_classes()，其实现如下：​​​​​​​

    def set_classes(self, text):        """Perform a forward pass with optional profiling, visualization, and embedding extraction."""        try:            import clip        except ImportError:            check_requirements("git+https://github.com/openai/CLIP.git")            import clip         model, _ = clip.load("ViT-B/32")        device = next(model.parameters()).device        text_token = clip.tokenize(text).to(device)        txt_feats = model.encode_text(text_token).to(dtype=torch.float32)        txt_feats = txt_feats / txt_feats.norm(p=2, dim=-1, keepdim=True)        self.txt_feats = txt_feats.reshape(-1, len(text), txt_feats.shape[-1])        self.model[-1].nc = len(text)

运行效果：

["person"]：

 ["person", "motorcycle"]：

运行过程中，时间消耗也是极小的，这也符合YOLO系列的风格。

Speed: 0.9ms preprocess, 43.5ms inference, 1.1ms postprocess per image at shape (1, 3, 448, 640)

可以预见的是，YOLO-World会是未来实时开放词汇检测领域的一个重量级选手。

THE END !

文章结束，感谢阅读。您的点赞，收藏，评论是我继续更新的动力。大家有推荐的公众号可以评论区留言，共同学习，一起进步。