44 lines
1.6 KiB
Python
44 lines
1.6 KiB
Python
import torch
|
||
import matplotlib.pyplot as plt # 新增：导入matplotlib用于保存图片
|
||
#################################### For Image ####################################
|
||
from PIL import Image
|
||
from sam3.model_builder import build_sam3_image_model
|
||
from sam3.model.sam3_image_processor import Sam3Processor
|
||
from sam3.visualization_utils import draw_box_on_image, normalize_bbox, plot_results
|
||
|
||
# Load the model
|
||
model = build_sam3_image_model()
|
||
processor = Sam3Processor(model)
|
||
|
||
# Load an image - 保留之前的RGB转换修复
|
||
image = Image.open("/home/quant/data/dev/sam3/assets/images/groceries.jpg").convert("RGB")
|
||
|
||
# 可选：打印图像信息，验证通道数
|
||
print(f"图像模式: {image.mode}, 尺寸: {image.size}")
|
||
|
||
# 处理图像
|
||
inference_state = processor.set_image(image)
|
||
|
||
# 文本提示推理
|
||
output = processor.set_text_prompt(state=inference_state, prompt="food")
|
||
|
||
# 获取推理结果
|
||
masks, boxes, scores = output["masks"], output["boxes"], output["scores"]
|
||
|
||
# 可视化并保存图片（核心修改部分）
|
||
# 1. 生成可视化结果
|
||
plot_results(image, inference_state)
|
||
# 2. 保存图片到当前目录，格式可选jpg/png，这里用jpg示例
|
||
plt.savefig("./sam3_food_detection_result.jpg", # 保存路径：当前目录，文件名自定义
|
||
dpi=150, # 图片分辨率，可选
|
||
bbox_inches='tight') # 去除图片周围空白
|
||
# 3. 关闭plt画布，避免内存占用
|
||
plt.close()
|
||
|
||
# 可选：打印输出信息
|
||
print(f"检测到的mask数量: {len(masks)}")
|
||
print(f"检测到的box数量: {len(boxes)}")
|
||
print(f"置信度分数: {scores}")
|
||
print("图片已保存到当前目录：./sam3_food_detection_result.jpg")
|
||
|