使用视觉语言模型(VLMs)进行目标检测

本文转自: 小白玩转Python,二旺


在过去,你必须自己训练模型,收集训练数据,但现在许多基础模型允许你在它们的基础上进行微调,以获得一个能够检测目标并与用户用自然语言互动的系统。有数百种模型和潜在应用场景,目标检测在这些场景中非常有用,尤其是随着小型语言模型的兴起,所以今天我们将尝试使用MLX上的Qwen2-VL-7B-Instruct-8bit。

我们将使用MLX-VLM,这是由Prince Canuma(Blaizzy)创建的一个包,他是一位热衷于开发和移植大型语言模型以兼容MLX的热情开发者,这个框架为我们用户抽象了很多代码,使我们能够用很少的代码行运行这些模型。现在让我们来看下面的代码片段。你会发现它非常简单。首先,你可以从Hugging Face定义模型,框架将下载所有相关组件。这个过程非常简单,因为这个库还提供了多个实用工具(apply_chat_template),可以将OpenAI的标准提示模板转换为小型VLMs所需的模板。

使用视觉语言模型(VLMs)进行目标检测

这里的一个重要注意事项是在编写代码时,这个库中的系统角色出现了一些问题,但未来很可能可以添加。但在本例中,我们在一个用户消息中传递任务和响应格式,基本上我们将要求模型识别所有对象并返回一个坐标列表,其中第一个顶部将是边界框的最小x/y坐标,后者将是最大坐标。同时,我们包括了对象名称,并要求模型以JSON对象的形式返回:

from mlx_vlm import load, apply_chat_template, generate
from mlx_vlm.utils import load_image


model, processor = load("mlx-community/Qwen2-VL-7B-Instruct-8bit")
config = model.config

image_path = "images/test.jpg"
image = load_image(image_path)

messages = [
    {
        "role": "user",
        "content": """detect all the objects in the image, return bounding boxes for all of them using the following format: [{
        "object": "object_name",
        "bboxes": [[xmin, ymin, xmax, ymax], [xmin, ymin, xmax, ymax], ...]
     }, ...]""",
    }
]
prompt = apply_chat_template(processor, config, messages)

output = generate(model, processor, image, prompt, max_tokens=1000, temperature=0.7)
print(output)

运行前面的代码后,你将收到一个JSON响应,正确识别了两辆卡车:

[{
    "object": "dump truck",
    "bboxes": [
        [100, 250, 380, 510]
    ]
}, {
    "object": "dump truck",
    "bboxes": [
        [550, 250, 830, 490]
    ]
}]

鉴于我们有了对象名称和边界框坐标,我们可以编写一个函数将这些结果绘制在图像上。代码如下:

import json
import re
import matplotlib.pyplot as plt
from PIL import Image, ImageDraw, ImageFont

def draw_and_plot_boxes_from_json(json_data, image_path):
    """
    Parses the JSON data to extract bounding box coordinates,
    scales them according to the image size, draws the boxes on the image,
    and plots the image.

    Args:
        json_data (str or list): The JSON data as a string or already parsed list.
        image_path (str): The path to the image file on which boxes are to be drawn.
    """
    # If json_data is a string, parse it into a Python object
    if isinstance(json_data, str):
        # Strip leading/trailing whitespaces
        json_data = json_data.strip()
        # Remove code fences if present
        json_data = re.sub(r"^```json\s*", "", json_data)
        json_data = re.sub(r"```$", "", json_data)
        json_data = json_data.strip()
        try:
            data = json.loads(json_data)
        except json.JSONDecodeError as e:
            print("Failed to parse JSON data:", e)
            print("JSON data was:", repr(json_data))
            return
    else:
        data = json_data

    # Open the image
    try:
        img = Image.open(image_path)
    except FileNotFoundError:
        print(f"Image file not found at {image_path}. Please check the path.")
        return

    draw = ImageDraw.Draw(img)
    width, height = img.size
    # Change this part for Windows OS
    # ImageFont.FreeTypeFont(r"C:\Windows\Fonts\CONSOLA.ttf", size=25)
    font = ImageFont.truetype("/System/Library/Fonts/Menlo.ttc", size=25)  # Process and draw boxes
    for item in data:
        object_type = item.get("object", "unknown")
        for bbox in item.get("bboxes", []):
            x1, y1, x2, y2 = bbox
            # Scale down coordinates from a 1000x1000 grid to the actual image size
            x1 = x1 * width / 1000
            y1 = y1 * height / 1000
            x2 = x2 * width / 1000
            y2 = y2 * height / 1000
            # Draw the rectangle on the image
            draw.rectangle([(x1, y1), (x2, y2)], outline="blue", width=5)
            text_position = (x1, y1)
            draw.text(text_position, object_type, fill="red", font=font)

    # Plot the image using matplotlib
    plt.figure(figsize=(8, 8))
    plt.imshow(img)
    plt.axis("off")  # Hide axes ticks
    plt.show()

绘制结果如下:

使用视觉语言模型(VLMs)进行目标检测


总结

VLMs正在快速发展。两年前,还没有能够适应MacBook并表现如此出色的模型。我个人的猜测是,这些模型将继续发展,最终达到像YOLO这样的模型的能力。还有很长的路要走,但正如你在这篇文章中看到的,设置这个演示非常容易。在边缘设备上开发这种应用的潜力是无限的,我相信它们将在采矿、石油和天然气、基础设施和监控等行业产生重大影响。最好的部分是我们甚至还没有讨论微调、RAG或提示工程,这只是模型能力的展示。


本文转自: 小白玩转Python,转载此文目的在于传递更多信息,版权归原作者所有。

最新文章