import cv2
import torch
import time
import numpy as np
import os
import sys# 检查文件是否存在
image_path = "npu2pm.JPG"
if not os.path.exists(image_path):print(f"错误：图像文件 '{image_path}' 不存在！")print("请确保图像文件在当前目录下")sys.exit(1)print("图像文件存在，继续执行...")# 设置设备
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"使用设备: {device}")# 性能优化配置
torch.backends.cudnn.benchmark = True
if device.type == "cuda":torch.backends.cuda.matmul.allow_tf32 = Truetorch.backends.cudnn.allow_tf32 = True# 图像和模型配置
mode_test = 'tiny'
scale = 1.0model_config = {"tiny": ("sam2.1_hiera_t.yaml", "sam2.1_hiera_tiny.pt"),"small": ("sam2.1_hiera_s.yaml", "sam2.1_hiera_small.pt"),"base": ("sam2.1_hiera_b.yaml", "sam2.1_hiera_base_plus.pt"),"large": ("sam2.1_hiera_l.yaml", "sam2.1_hiera_large.pt")
}model_type, model_path = model_config[mode_test]
checkpoint = f"../checkpoints/{model_path}"
model_cfg = f"../sam2/configs/sam2.1/{model_type}"print(f"模型配置: {model_cfg}")
print(f"检查点: {checkpoint}")# 检查模型文件是否存在
if not os.path.exists(checkpoint.replace("../checkpoints/", "")) and not os.path.exists(checkpoint):print(f"警告：模型文件可能不存在于: {checkpoint}")# 加载并预处理图像
print("正在加载图像...")
image_cv = cv2.imread(image_path)
if image_cv is None:raise ValueError("无法加载图像！")height, width = image_cv.shape[:2]
print(f"原始图像尺寸: {width}x{height}")new_width = int(width * scale)
new_height = int(height * scale)image = cv2.resize(image_cv, (new_width, new_height))
image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)print(f"调整后图像尺寸: {new_width}x{new_height}")# 构建模型
print("正在加载SAM2模型...")
try:from sam2.build_sam import build_sam2from sam2.sam2_image_predictor import SAM2ImagePredictorstart_load = time.time()sam2 = build_sam2(model_cfg, checkpoint, device=device)predictor = SAM2ImagePredictor(sam2)end_load = time.time()print(f"模型加载完成，耗时: {end_load - start_load:.2f} 秒")except ImportError as e:print(f"导入错误: {e}")print("请确保sam2模块在Python路径中")sys.exit(1)
except Exception as e:print(f"模型加载错误: {e}")print("创建模拟模式进行测试...")# 创建模拟模式进行测试class MockPredictor:def set_image(self, image):print("模拟: set_image调用")def predict(self, **kwargs):print("模拟: predict调用")# 返回模拟的maskmask = np.zeros((new_height, new_width), dtype=bool)h, w = mask.shape# 创建一个简单的矩形maskmask[h//4:3*h//4, w//4:3*w//4] = Truereturn [mask], [0.95], Nonepredictor = MockPredictor()print("使用模拟模式进行测试")# 清理GPU缓存
if device.type == "cuda":torch.cuda.empty_cache()# 交互式选择框
class BoxSelector:def __init__(self, image):self.image = imageself.drawing = Falseself.box = []self.temp_image = image.copy()self.result_image = Noneself.mask_window_name = "Segmentation Mask"self.result_window_name = "Segmentation Result"# 创建显示窗口cv2.namedWindow("Select Box", cv2.WINDOW_NORMAL)cv2.resizeWindow("Select Box", 800, 600)cv2.namedWindow(self.mask_window_name, cv2.WINDOW_NORMAL)cv2.resizeWindow(self.mask_window_name, 800, 600)cv2.namedWindow(self.result_window_name, cv2.WINDOW_NORMAL)cv2.resizeWindow(self.result_window_name, 800, 600)cv2.setMouseCallback("Select Box", self.draw_box)def draw_box(self, event, x, y, flags, param):if event == cv2.EVENT_LBUTTONDOWN:print(f"鼠标按下: ({x}, {y})")self.drawing = Trueself.box = [x, y, x, y]self.update_display()elif event == cv2.EVENT_MOUSEMOVE:if self.drawing:self.box[2:] = [x, y]self.update_display()elif event == cv2.EVENT_LBUTTONUP:print(f"鼠标释放: ({x}, {y})")self.drawing = Falseself.box[2:] = [x, y]# 确保x_min < x_max, y_min < y_maxself.box = [min(self.box[0], self.box[2]), min(self.box[1], self.box[3]), max(self.box[0], self.box[2]), max(self.box[1], self.box[3])]# 检查框的有效性box_width = self.box[2] - self.box[0]box_height = self.box[3] - self.box[1]print(f"选择的框: {self.box}, 尺寸: {box_width}x{box_height}")if box_width < 5 or box_height < 5:print("框太小，请重新选择")self.temp_image = self.image.copy()cv2.imshow("Select Box", self.temp_image)returnself.update_display()print("开始分割处理...")# 执行分割self.perform_segmentation(self.box)def update_display(self):"""更新显示图像"""self.temp_image = self.image.copy()if len(self.box) == 4:cv2.rectangle(self.temp_image, (self.box[0], self.box[1]), (self.box[2], self.box[3]), (0, 255, 0), 2)cv2.imshow("Select Box", self.temp_image)def perform_segmentation(self, box):print("执行分割...")start_time = time.time()try:predictor.set_image(image_rgb)masks, scores, logits = predictor.predict(point_coords=None,point_labels=None,box=np.array(box),multimask_output=False)end_time = time.time()processing_time = end_time - start_timeprint(f"分割完成，耗时: {processing_time:.2f} 秒")print(f"生成掩码数量: {len(masks)}")if len(masks) > 0:# 创建基础图像base_image = self.image.copy()cv2.rectangle(base_image, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2)# 创建纯mask显示图像mask_display = np.zeros_like(base_image)for i, mask in enumerate(masks):print(f"处理掩码 {i+1}, 分数: {scores[i]:.3f}")if hasattr(mask, 'cpu'):mask = mask.cpu().numpy()if mask.dtype != bool:mask = mask.astype(bool)# 生成鲜艳的颜色color = [np.random.randint(150, 256),np.random.randint(150, 256),np.random.randint(150, 256)]# 增强mask显示效果mask_alpha = 0.7 # mask透明度border_width = 2 # 边界宽度# 1. 创建纯mask显示mask_only = np.zeros_like(base_image)for c in range(3):mask_only[:, :, c][mask] = color[c]cv2.imshow(self.mask_window_name, mask_only)# 2. 创建结果图像（带原图）result_image = base_image.copy()# 创建彩色掩码colored_mask = np.zeros_like(result_image)for c in range(3):colored_mask[:, :, c][mask] = color[c]# 混合图像（增强mask效果）result_image = cv2.addWeighted(result_image, 1 - mask_alpha, colored_mask, mask_alpha, 0)# 添加边界contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)cv2.drawContours(result_image, contours, -1, color, border_width)# 显示分数y_coords, x_coords = np.where(mask)if len(x_coords) > 0:center_x, center_y = np.mean(x_coords), np.mean(y_coords)cv2.putText(result_image, f"Score: {scores[i]:.3f}", (int(center_x), int(center_y)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)# 显示结果cv2.imshow(self.result_window_name, result_image)self.result_image = result_image# 保存结果result_path = "box_result_interactive.jpg"cv2.imwrite(result_path, self.result_image)print(f"结果已保存: {result_path}")else:print("未生成掩码")# 显示空白结果cv2.imshow(self.mask_window_name, np.zeros_like(self.image))cv2.imshow(self.result_window_name, self.image.copy())except Exception as e:print(f"分割错误: {e}")import tracebacktraceback.print_exc()# 显示错误信息error_img = np.zeros_like(self.image)cv2.putText(error_img, f"Error: {str(e)}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2)cv2.imshow(self.result_window_name, error_img)# 创建选择器实例
selector = BoxSelector(image)print("\n=== 交互式分割程序 ===")
print("操作说明:")
print("1. 在'Select Box'窗口上按住鼠标左键并拖动选择框")
print("2. 释放鼠标后自动执行分割")
print("3. 'Segmentation Mask'窗口显示纯mask效果")
print("4. 'Segmentation Result'窗口显示混合结果")
print("5. 按 'r' 键重置选择")
print("6. 按 ESC 或 'q' 键退出")
print("=" + "="*30)# 初始显示
cv2.imshow("Select Box", image)
cv2.imshow(selector.mask_window_name, np.zeros_like(image))
cv2.imshow(selector.result_window_name, image.copy())
print("窗口已显示，等待用户输入...")# 主循环
try:while True:key = cv2.waitKey(100) & 0xFF # 100ms延迟，减少CPU使用if key == 27 or key == ord('q'): # ESC或q键print("用户请求退出")breakelif key == ord('r'): # 重置键print("重置选择")selector.box = []selector.temp_image = image.copy()cv2.imshow("Select Box", selector.temp_image)cv2.imshow(selector.mask_window_name, np.zeros_like(image))cv2.imshow(selector.result_window_name, image.copy())except KeyboardInterrupt:print("程序被中断")finally:cv2.destroyAllWindows()print("程序结束")# 清理内存
if device.type == "cuda":torch.cuda.empty_cache()
