78 lines
2.9 KiB
Python
78 lines
2.9 KiB
Python
# 导入必要的库
|
||
import os # 操作系统模块
|
||
import torch # PyTorch深度学习框架
|
||
from torch.utils.data import DataLoader # 数据加载器
|
||
from torchvision import transforms # 图像变换工具
|
||
from PIL import Image # 图像处理库
|
||
import cv2 # OpenCV库
|
||
import numpy as np # NumPy库
|
||
|
||
# 导入自定义模块
|
||
from unet_coco_segmentation import CocoSegDataset # COCO数据集类
|
||
import segmentation_models_pytorch as smp # 分割模型库
|
||
|
||
# 配置参数
|
||
TEST_DIR = '../data/test' # 测试集目录
|
||
TEST_ANN = '../data/test/_annotations.coco.json' # 测试集标注文件
|
||
MODEL_PATH = 'unet_coco_segmentation.pth' # 模型权重路径
|
||
OUTPUT_DIR = 'output/unet_results' # 输出目录
|
||
|
||
# 创建输出目录（如果不存在）
|
||
os.makedirs(OUTPUT_DIR, exist_ok=True)
|
||
|
||
# 设备配置（GPU或CPU）
|
||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||
|
||
# 数据预处理配置（与训练时保持一致）
|
||
transform = transforms.Compose([
|
||
transforms.Resize((256, 256)), # 调整尺寸
|
||
transforms.ToTensor(), # 转换为张量
|
||
transforms.Normalize(mean=(0.485,0.456,0.406), std=(0.229,0.224,0.225)), # 归一化
|
||
])
|
||
|
||
# 加载测试集
|
||
test_dataset = CocoSegDataset(
|
||
root_dir=TEST_DIR,
|
||
annotation_file=TEST_ANN,
|
||
transforms=None, # 在Dataset里单独处理
|
||
mask_transforms=None
|
||
)
|
||
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
|
||
|
||
# 初始化模型（与训练时相同）
|
||
model = smp.Unet(
|
||
encoder_name='resnet34', # 编码器名称
|
||
encoder_weights=None, # 不使用预训练权重
|
||
in_channels=3, # 输入通道数
|
||
classes=1 # 输出类别数
|
||
)
|
||
# 加载预训练模型权重
|
||
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
|
||
model.to(device) # 将模型移动到指定设备
|
||
model.eval() # 设置为评估模式
|
||
print("成功加载模型权重并切换到评估模式")
|
||
|
||
# 遍历测试集
|
||
for img, mask_true in test_loader:
|
||
# 图像保存为Tensor，batch=1
|
||
img = img.to(device)
|
||
# 获取图像ID
|
||
img_id = test_dataset.image_ids[test_loader.dataset.image_ids.index(test_dataset.image_ids[0])]
|
||
# 进行预测
|
||
with torch.no_grad():
|
||
output = model(img)
|
||
output_prob = torch.sigmoid(output).squeeze().cpu().numpy()
|
||
# 获取原始图像尺寸
|
||
img_info = next(item for item in test_dataset.coco['images'] if item['id']==test_dataset.image_ids[0])
|
||
orig_w, orig_h = img_info['width'], img_info['height']
|
||
# 调整输出尺寸到原始图像大小
|
||
output_prob = cv2.resize(output_prob, (orig_w, orig_h), interpolation=cv2.INTER_LINEAR)
|
||
|
||
# 二值化处理
|
||
threshold = 0.5
|
||
output_mask = (output_prob > threshold).astype(np.uint8) * 255
|
||
|
||
# 保存结果
|
||
output_path = os.path.join(OUTPUT_DIR, f"{img_id}_mask.png")
|
||
cv2.imwrite(output_path, output_mask)
|
||
print(f"Saved mask for image {img_id} to {output_path}") |