""" 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()