import time
import psutil
import cv2
from segment_anything import sam_model_registry, SamPredictor
# 模拟图像数据
image_paths = ['notebooks/images/dog.jpg', 'notebooks/images/groceries.jpg', 'notebooks/images/truck.jpg']
total_prediction_time = 0.0
# 加载SAM预训练模型
sam_checkpoint = "./models/sam_vit_b_01ec64.pth"
model_type = "vit_b"
device = "cpu" # 或者 "cuda",取决于你的设备
sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
sam.to(device=device)
predictor = SamPredictor(sam)
# 对每张图片进行预测
for image_path in image_paths:
start_time = time.time()
# 读取图像
image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
# 设置待预测的图像
predictor.set_image(image)
# 执行预测操作
masks, scores, logits = predictor.predict(
multimask_output=True,
)
end_time = time.time()
# 计算每张图片的预测时间
prediction_time = end_time - start_time
total_prediction_time += prediction_time
print(f"Prediction time for {image_path}: {prediction_time} seconds")
# 计算平均预测时间
average_prediction_time = total_prediction_time / len(image_paths)
print(f"Average prediction time: {average_prediction_time} seconds")
# 获取系统资源占用情况
memory_percent = psutil.virtual_memory().percent
cpu_percent = psutil.cpu_percent(interval=1)
print(f"Memory percent: {memory_percent}%")
print(f"CPU percent: {cpu_percent}%")
python评估图像边缘检测模型性能--segment_anything模型
于 2024-03-09 16:59:40 首次发布