Mask R-CNN instance segmentation training code
This commit is contained in:
flk
+262
@@ -0,0 +1,262 @@
|
||||
"""
|
||||
Mask R-CNN 实例分割推理与可视化脚本
|
||||
|
||||
用法:
|
||||
# 对单张图片进行推理
|
||||
python predict.py --image dataset/1.jpg --checkpoint checkpoints/mask_rcnn_best.pth
|
||||
|
||||
# 对整个数据集目录进行推理
|
||||
python predict.py --dataset dataset --checkpoint checkpoints/mask_rcnn_best.pth
|
||||
|
||||
# 调整置信度阈值
|
||||
python predict.py --image dataset/1.jpg --checkpoint checkpoints/mask_rcnn_best.pth --threshold 0.7
|
||||
|
||||
# 同时显示 ground truth 对比
|
||||
python predict.py --image dataset/1.jpg --checkpoint checkpoints/mask_rcnn_best.pth --show-gt
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
from typing import Dict, Optional
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import torch
|
||||
|
||||
from utils import (
|
||||
get_label_map,
|
||||
get_image_json_pairs,
|
||||
shapes_to_masks,
|
||||
visualize_instances,
|
||||
)
|
||||
from dataset import get_val_transforms
|
||||
from model import get_model
|
||||
|
||||
|
||||
def load_checkpoint(checkpoint_path: str, device: torch.device):
|
||||
"""
|
||||
加载训练好的模型检查点。
|
||||
|
||||
Args:
|
||||
checkpoint_path: 检查点文件路径
|
||||
device: 计算设备
|
||||
|
||||
Returns:
|
||||
model: 加载好权重的模型
|
||||
label_map: 标签名→类别ID 映射
|
||||
"""
|
||||
checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
||||
label_map = checkpoint.get("label_map", {})
|
||||
|
||||
num_classes = max(label_map.values()) + 1 # +1 for background
|
||||
|
||||
model = get_model(num_classes=num_classes, pretrained=False)
|
||||
model.load_state_dict(checkpoint["model_state_dict"])
|
||||
model = model.to(device)
|
||||
model.eval()
|
||||
|
||||
# 构建 id→name 映射
|
||||
id_to_name = {v: k for k, v in label_map.items()}
|
||||
|
||||
return model, label_map, id_to_name
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def predict_single(
|
||||
model: torch.nn.Module,
|
||||
image: np.ndarray,
|
||||
device: torch.device,
|
||||
threshold: float = 0.5,
|
||||
) -> Dict:
|
||||
"""
|
||||
对单张图像进行推理。
|
||||
|
||||
Args:
|
||||
model: 模型
|
||||
image: [H, W, 3] RGB uint8 图像
|
||||
device: 计算设备
|
||||
threshold: 置信度阈值
|
||||
|
||||
Returns:
|
||||
包含 boxes, masks, labels, scores 的字典
|
||||
"""
|
||||
# 预处理
|
||||
image_float = image.astype(np.float32) / 255.0
|
||||
image_tensor = torch.as_tensor(image_float, dtype=torch.float32)
|
||||
image_tensor = image_tensor.permute(2, 0, 1).unsqueeze(0).to(device)
|
||||
|
||||
# 推理
|
||||
outputs = model(image_tensor)
|
||||
output = outputs[0]
|
||||
|
||||
# 过滤低置信度
|
||||
keep = output["scores"] > threshold
|
||||
boxes = output["boxes"][keep].cpu().numpy()
|
||||
masks = output["masks"][keep].cpu().numpy()
|
||||
labels = output["labels"][keep].cpu().numpy()
|
||||
scores = output["scores"][keep].cpu().numpy()
|
||||
|
||||
# mask 阈值化（> 0.5）
|
||||
masks = (masks[:, 0, :, :] > 0.5).astype(np.uint8)
|
||||
|
||||
return {
|
||||
"boxes": boxes,
|
||||
"masks": masks,
|
||||
"labels": labels,
|
||||
"scores": scores,
|
||||
}
|
||||
|
||||
|
||||
def visualize_and_save(
|
||||
image: np.ndarray,
|
||||
result: Dict,
|
||||
id_to_name: Dict[int, str],
|
||||
output_path: str,
|
||||
threshold: float = 0.5,
|
||||
):
|
||||
"""
|
||||
可视化预测结果并保存图像。
|
||||
|
||||
Args:
|
||||
image: 原始图像 [H, W, 3] RGB
|
||||
result: 预测结果字典
|
||||
id_to_name: 类别ID→名称映射
|
||||
output_path: 输出路径
|
||||
threshold: 置信度阈值
|
||||
"""
|
||||
vis_img = visualize_instances(
|
||||
image,
|
||||
result["boxes"],
|
||||
result["masks"],
|
||||
result["labels"],
|
||||
result["scores"],
|
||||
label_map=id_to_name,
|
||||
score_threshold=threshold,
|
||||
)
|
||||
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
Image.fromarray(vis_img).save(output_path)
|
||||
print(f" -> 结果已保存: {output_path}")
|
||||
|
||||
|
||||
def visualize_gt(
|
||||
image: np.ndarray,
|
||||
json_path: str,
|
||||
label_map: Dict[str, int],
|
||||
id_to_name: Dict[int, str],
|
||||
output_path: str,
|
||||
):
|
||||
"""可视化 ground truth 标注"""
|
||||
with open(json_path, "r", encoding="utf-8") as f:
|
||||
anno = json.load(f)
|
||||
|
||||
height = anno.get("imageHeight", image.shape[0])
|
||||
width = anno.get("imageWidth", image.shape[1])
|
||||
|
||||
masks, labels, boxes = shapes_to_masks(
|
||||
anno.get("shapes", []), height, width, label_map
|
||||
)
|
||||
|
||||
if len(labels) == 0:
|
||||
print(" -> 无标注实例")
|
||||
return
|
||||
|
||||
vis_img = visualize_instances(
|
||||
image, boxes, masks, labels, None, id_to_name, 0.0
|
||||
)
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
Image.fromarray(vis_img).save(output_path)
|
||||
print(f" -> Ground Truth 已保存: {output_path}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Mask R-CNN 实例分割推理")
|
||||
parser.add_argument("--image", type=str, default=None,
|
||||
help="单张图像路径")
|
||||
parser.add_argument("--dataset", type=str, default=None,
|
||||
help="数据集目录路径（批量推理）")
|
||||
parser.add_argument("--checkpoint", type=str, default="checkpoints/mask_rcnn_best.pth",
|
||||
help="模型检查点路径 (默认: checkpoints/mask_rcnn_best.pth)")
|
||||
parser.add_argument("--output-dir", type=str, default="predictions",
|
||||
help="输出目录 (默认: predictions)")
|
||||
parser.add_argument("--threshold", type=float, default=0.5,
|
||||
help="置信度阈值 (默认: 0.5)")
|
||||
parser.add_argument("--show-gt", action="store_true",
|
||||
help="同时可视化 ground truth 对比")
|
||||
parser.add_argument("--min-size", type=int, default=800,
|
||||
help="输入图像最小边尺寸 (默认: 800)")
|
||||
parser.add_argument("--max-size", type=int, default=1333,
|
||||
help="输入图像最大边尺寸 (默认: 1333)")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# 设置设备
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"使用设备: {device}")
|
||||
|
||||
# 检查检查点
|
||||
ckpt_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.checkpoint)
|
||||
if not os.path.exists(ckpt_path):
|
||||
print(f"错误: 检查点文件不存在: {ckpt_path}")
|
||||
print("请先运行 train.py 训练模型，或使用 --checkpoint 指定正确的路径。")
|
||||
sys.exit(1)
|
||||
|
||||
# 加载模型
|
||||
model, label_map, id_to_name = load_checkpoint(ckpt_path, device)
|
||||
print(f"模型已加载 | 标签映射: {label_map}")
|
||||
print(f"置信度阈值: {args.threshold}")
|
||||
|
||||
output_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.output_dir)
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
# 收集要推理的图像
|
||||
if args.image:
|
||||
image_paths = [args.image]
|
||||
json_paths = [args.image.rsplit(".", 1)[0] + ".json"]
|
||||
elif args.dataset:
|
||||
dataset_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), args.dataset)
|
||||
pairs = get_image_json_pairs(dataset_path)
|
||||
image_paths = [p[0] for p in pairs]
|
||||
json_paths = [p[1] for p in pairs]
|
||||
else:
|
||||
print("错误: 请指定 --image 或 --dataset")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"\n共 {len(image_paths)} 张图像待推理\n")
|
||||
|
||||
for idx, (image_path, json_path) in enumerate(zip(image_paths, json_paths)):
|
||||
if not os.path.exists(image_path):
|
||||
print(f"[{idx+1}] 跳过（图像不存在）: {image_path}")
|
||||
continue
|
||||
|
||||
print(f"[{idx+1}/{len(image_paths)}] {os.path.basename(image_path)}")
|
||||
|
||||
# 读取图像
|
||||
image = np.array(Image.open(image_path).convert("RGB"))
|
||||
|
||||
# 推理
|
||||
result = predict_single(model, image, device, args.threshold)
|
||||
|
||||
num_detections = len(result["scores"])
|
||||
print(f" -> 检测到 {num_detections} 个实例")
|
||||
for i in range(num_detections):
|
||||
name = id_to_name.get(int(result["labels"][i]), str(result["labels"][i]))
|
||||
print(f" {name}: {result['scores'][i]:.3f}")
|
||||
|
||||
# 可视化并保存
|
||||
base_name = os.path.splitext(os.path.basename(image_path))[0]
|
||||
output_path = os.path.join(output_dir, f"{base_name}_pred.png")
|
||||
visualize_and_save(image, result, id_to_name, output_path, args.threshold)
|
||||
|
||||
# 可选：可视化 ground truth
|
||||
if args.show_gt and os.path.exists(json_path):
|
||||
gt_output_path = os.path.join(output_dir, f"{base_name}_gt.png")
|
||||
visualize_gt(image, json_path, label_map, id_to_name, gt_output_path)
|
||||
|
||||
print(f"\n推理完成！结果保存在: {output_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user