【Segment Anything Model】十四:原始SAM模型如何传入多框

之前第二三篇有更新过单点,多点,单框。本篇加上多框输入。

先确定一下目录

新建test_boxes.py文件,复制以下代码

python 复制代码
import sys
import torch
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
from Net.segment_anything import SamPredictor,sam_model_registry
import cv2

image = cv2.imread('1.jpg')  # 读取的图像以NumPy数组的形式存储在变量image中
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)  # 将图像从BGR颜色空间转换为RGB颜色空间,还原图片色彩(图像处理库所认同的格式)

model_type = "vit_b"  # 定义模型类型
device = "cuda"  # "cpu"  or  "cuda"
sam = sam_model_registry[model_type](checkpoint='sam_vit_b_01ec64.pth')
sam.to(device=device)  # 定义模型参数
predictor = SamPredictor(sam)  # 调用预测模型

predictor.set_image(image)

def show_mask(mask, ax, random_color=False):
    if random_color:
        color = np.concatenate([np.random.random(3), np.array([0.6])], axis=0)
    else:
        color = np.array([30 / 255, 144 / 255, 255 / 255, 0.6])
    h, w = mask.shape[-2:]
    mask_image = mask.reshape(h, w, 1) * color.reshape(1, 1, -1)
    ax.imshow(mask_image)


def show_points(coords, labels, ax, marker_size=375):
    pos_points = coords[labels == 1]
    neg_points = coords[labels == 0]
    ax.scatter(pos_points[:, 0], pos_points[:, 1], color='green', marker='*', s=marker_size, edgecolor='white',
               linewidth=1.25)
    ax.scatter(neg_points[:, 0], neg_points[:, 1], color='red', marker='*', s=marker_size, edgecolor='white',
               linewidth=1.25)


def show_box(box, ax):
    x0, y0 = box[0], box[1]
    w, h = box[2] - box[0], box[3] - box[1]
    ax.add_patch(plt.Rectangle((x0, y0), w, h, edgecolor='green', facecolor=(0, 0, 0, 0), lw=2))

# 多框
# input_boxes = torch.tensor([
#     [75, 275, 1725, 850],
#     [425, 600, 700, 875],
#     [1375, 550, 1650, 800],
#     [1240, 675, 1400, 750],
# ], device=predictor.device)  # 假设为目标检测的预测结果

# 单框
input_boxes = torch.tensor([
    [75, 275, 1725, 850]
], device=predictor.device)  # 假设为目标检测的预测结果
input_boxes = input_boxes / 2
transformed_boxes = predictor.transform.apply_boxes_torch(input_boxes, image.shape[:2])
masks, _, _ = predictor.predict_torch(
    point_coords=None,
    point_labels=None,
    boxes=transformed_boxes,
    multimask_output=False,
)


plt.figure(figsize=(10, 10))
plt.imshow(image)
for mask in masks:
    show_mask(mask.cpu().numpy(), plt.gca(), random_color=True)
for box in input_boxes:
    show_box(box.cpu().numpy(), plt.gca())
plt.axis('off')
plt.show()

打断点

在如下位置打断点,debug运行,观察单框和多框传入有什么维度不同,自己训练的时候按照对应维度组装。

相关推荐
橡晟3 小时前
深度学习入门:让神经网络变得“深不可测“⚡(二)
人工智能·python·深度学习·机器学习·计算机视觉
墨尘游子3 小时前
神经网络的层与块
人工智能·python·深度学习·机器学习
倔强青铜33 小时前
苦练Python第18天:Python异常处理锦囊
开发语言·python
企鹅与蟒蛇4 小时前
Ubuntu-25.04 Wayland桌面环境安装Anaconda3之后无法启动anaconda-navigator问题解决
linux·运维·python·ubuntu·anaconda
autobaba4 小时前
编写bat文件自动打开chrome浏览器,并通过selenium抓取浏览器操作chrome
chrome·python·selenium·rpa
Rvelamen5 小时前
LLM-SECURITY-PROMPTS大模型提示词攻击测评基准
人工智能·python·安全
【本人】5 小时前
Django基础(一)———创建与启动
后端·python·django
SHIPKING3936 小时前
【python】基于pygame实现动态粒子爱心
开发语言·python·pygame
kk_stoper8 小时前
如何通过API查询实时能源期货价格
java·开发语言·javascript·数据结构·python·能源
java1234_小锋8 小时前
【NLP舆情分析】基于python微博舆情分析可视化系统(flask+pandas+echarts) 视频教程 - 架构搭建
python·自然语言处理·flask