【扒代码】demo.py

def on_click(event):
    # 当鼠标按下时，记录边界框的起始点
    # 使用了全局变量ix和iy来存储起始点的x和y坐标
    global ix, iy  # 声明ix和iy为全局变量，以便在函数外部访问和修改它们的值
    ix, iy = event.xdata, event.ydata  # 从事件对象event中获取鼠标点击的x和y坐标

    # 连接鼠标释放事件到on_release函数
    # 当鼠标按钮释放时，将调用on_release函数来完成边界框的绘制
    fig.canvas.mpl_connect('button_release_event', on_release)  # fig是matplotlib的图形对象，canvas是其画布对象

功能解释：

• on_click函数是一个事件处理函数，它的作用是在用户点击鼠标按钮时被调用。
• 函数内部使用了global关键字来声明ix和iy为全局变量。这是因为我们需要在函数外部也能访问和修改这两个变量的值，以便在鼠标释放时能够记录边界框的起始点。
• event.xdata和event.ydata分别代表了鼠标事件在画布上的x和y坐标。这些值是从触发事件的对象event中获取的，它包含了鼠标点击事件的相关信息。
• fig.canvas.mpl_connect是一个连接matplotlib画布事件的函数。在这里，它被用来连接'button_release_event'事件，即鼠标按钮释放的事件。当用户释放鼠标按钮时，将调用on_release函数。
• 通过这种方式，on_click函数记录了用户开始绘制边界框的位置，而当用户释放鼠标按钮时，on_release函数将被触发来完成边界框的绘制，并存储边界框的坐标。

这种事件处理机制允许用户通过交互式地点击和拖动鼠标来在图像上绘制边界框，这在图像标注、对象检测和图像分割等任务中非常有用。

def on_release(event):
    # 当鼠标释放时，记录边界框的结束点
    global ix, iy  # 声明ix和iy为全局变量，以便访问在on_click函数中设置的值
    x, y = event.xdata, event.ydata  # 从事件对象event中获取鼠标释放时的x和y坐标

    # 计算边界框的宽度和高度
    width = x - ix  # 宽度是结束点x坐标减去起始点x坐标
    height = y - iy  # 高度是结束点y坐标减去起始点y坐标

    # 创建一个矩形补丁（边界框）并添加到轴对象上
    # patches.Rectangle用于在图像上绘制矩形区域
    rect = patches.Rectangle(  # 创建矩形补丁的实例
        (ix, iy),  # 矩形的左上角坐标（起始点）
        width, height,  # 矩形的宽度和高度
        edgecolor='r',  # 矩形边界的颜色，这里是红色
        facecolor='none'  # 矩形内部的颜色，这里是透明
    )
    ax.add_patch(rect)  # 将矩形补丁添加到当前的轴对象ax上

    # 存储边界框的坐标到列表中
    # 边界框的坐标以(左上角x, 左上角y, 右下角x, 右下角y)的形式存储
    bounding_boxes.append((ix, iy, ix + width, iy + height))

    # 调用plt.draw()来更新绘图窗口，显示新绘制的边界框
    plt.draw()

功能解释：

• on_release函数是一个事件处理函数，用于处理用户在matplotlib绘图窗口中释放鼠标按钮的动作。
• 函数开始时，使用global关键字声明ix和iy为全局变量，这样函数就可以访问在on_click函数中记录的鼠标点击起始点的坐标。
• 通过从event对象获取的xdata和ydata属性，函数捕捉到鼠标释放时的位置，即边界框的结束点。
• 根据起始点和结束点的坐标，函数计算出边界框的宽度和高度。
• 使用patches.Rectangle创建一个矩形补丁，该补丁表示边界框，并设置其左上角坐标、宽度、高度和颜色。
• 通过调用ax.add_patch(rect)将创建的矩形补丁添加到当前的轴对象ax上，从而在图像上显示边界框。
• 将边界框的坐标以元组的形式添加到全局列表bounding_boxes中，该列表用于存储所有绘制的边界框的坐标。
• 最后，调用plt.draw()函数来重新绘制图像，使得新绘制的边界框能够显示在matplotlib的绘图窗口中。

@torch.no_grad()  # 装饰器，表示此函数在执行时不会追踪梯度，常用于推理阶段
def demo(args):
    # 定义全局变量和函数参数
    img_path = "material//458.jpg"  # 图像路径
    global fig, ax  # 声明全局变量fig和ax，用于matplotlib绘图

    gpu = 0  # 指定使用的GPU编号
    # torch.cuda.set_device(gpu)  # 指定使用的GPU设备，被注释掉表示不使用GPU
    # device = torch.device('cpu')  # 设备设置，被注释掉表示使用CPU

    # 构建模型并使用DataParallel包装，用于多GPU训练
    model = DataParallel(
        build_model(args)  # 根据参数args构建模型
        # device_ids=[gpu],  # 指定使用的GPU编号列表，被注释掉
        # output_device=gpu  # 指定输出设备的编号，被注释掉
    )

    # 加载预训练模型权重，使用CPU映射
    model.load_state_dict(
        torch.load(os.path.join(args.model_path, 'DAVE_3_shot.pth'),
                   map_location=torch.device('cpu'),
                   weights_only=True)['model'],
        strict=False  # 如果模型定义和权重不完全匹配，不抛出错误
    )

    # 加载特征比较模块的预训练权重
    pretrained_dict_feat = {
        k.split("feat_comp.")[1]: v for k, v in
        torch.load(os.path.join(args.model_path, 'verification.pth'),
                   map_location=torch.device('cpu')
                   )['model'].items() if 'feat_comp' in k
    }
    model.module.feat_comp.load_state_dict(pretrained_dict_feat)  # 加载权重到特征比较模块
    model.eval()  # 将模型设置为评估模式

    # 打开图像并显示，允许用户通过鼠标点击绘制边界框
    image = Image.open(img_path).convert("RGB")  # 打开图像并转换为RGB模式
    fig, ax = plt.subplots(1)  # 创建一个图形和轴
    ax.imshow(image)  # 在轴上显示图像
    cid = fig.canvas.mpl_connect('button_press_event', on_click)  # 连接鼠标点击事件到on_click函数
    plt.title("Click and drag to draw bboxes, then close window")  # 设置图形标题
    plt.show()  # 显示图像

    # 将用户绘制的边界框转换为张量
    bboxes = torch.tensor(bounding_boxes)  # 将边界框列表转换为张量

    # 对图像和边界框进行尺寸调整
    img, bboxes, scale = resize(image, bboxes)
    img = img.unsqueeze(0)  # 增加一个批次维度
    bboxes = bboxes.unsqueeze(0)  # 增加一个批次维度

    # 使用模型进行预测
    density_map, _, _, predicted_bboxes = model(img, bboxes=bboxes)  # 模型预测

    # 显示预测结果
    plt.clf()  # 清除当前图形
    plt.imshow(image)  # 重新显示图像
    pred_boxes = predicted_bboxes.box.cpu() / torch.tensor([scale[0], scale[1], scale[0], scale[1]])  # 预测的边界框缩放回原始尺寸
    for i in range(len(pred_boxes)):
        box = pred_boxes[i]  # 获取预测的边界框
        plt.plot([box[0], box[0], box[2], box[2], box[0]], [box[1], box[3], box[3], box[1], box[1]], linewidth=2, color='red')  # 在图像上绘制边界框
    plt.title("Dmap count:" + str(round(density_map.sum().item(), 1)) + " Box count:" + str(len(pred_boxes)))  # 设置标题显示预测的密度图计数和边界框计数
    plt.show()  # 显示预测结果

• demo函数是一个演示函数，用于展示如何使用DAVE模型进行对象检测和计数。
• 函数首先设置模型的路径和参数，然后加载预训练的权重。
• 使用matplotlib创建一个交互式窗口，允许用户通过鼠标点击来绘制边界框。
• 用户绘制完边界框后，这些边界框被转换为张量，并进行尺寸调整以匹配模型的输入要求。
• 调整后的图像和边界框作为输入传递给模型进行预测。
• 最后，函数显示模型预测的密度图和边界框，并在原始图像上绘制这些边界框