SAM2 图像分割（3）鼠标选择多框 摄像头实时分割显示

 

import cv2
import torch
import time
import numpy as np
import os
import sysimport sys
sys.path.append('/home/r9000k/v2_project/v5_samyolo/2分割/sam2-main/test')# 设置设备
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}")# 构建模型
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((480, 640), 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 CameraBoxSelector:def __init__(self):self.cap = cv2.VideoCapture(0)  # 打开默认摄像头if not self.cap.isOpened():raise RuntimeError("无法打开摄像头！")# 设置摄像头分辨率self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)self.drawing = Falseself.current_box = []self.boxes = []  # 存储所有框self.current_frame = Noneself.result_image = Noneself.mask_window_name = "Segmentation Mask"self.result_window_name = "Segmentation Result"self.need_update = Trueself.monitoring_mode = False  # 是否进入实时监测模式self.last_process_time = 0self.process_interval = 0.2  # 处理间隔(秒)# 创建显示窗口cv2.namedWindow("Camera Feed", cv2.WINDOW_NORMAL)cv2.resizeWindow("Camera Feed", 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("Camera Feed", self.draw_box)def draw_box(self, event, x, y, flags, param):if self.monitoring_mode:return  # 监测模式下不响应鼠标事件if event == cv2.EVENT_LBUTTONDOWN:print(f"鼠标按下: ({x}, {y})")self.drawing = Trueself.current_box = [x, y, x, y]self.need_update = Trueelif event == cv2.EVENT_MOUSEMOVE:if self.drawing:self.current_box[2:] = [x, y]self.need_update = Trueelif event == cv2.EVENT_LBUTTONUP:print(f"鼠标释放: ({x}, {y})")self.drawing = Falseself.current_box[2:] = [x, y]# 确保x_min < x_max, y_min < y_maxself.current_box = [min(self.current_box[0], self.current_box[2]), min(self.current_box[1], self.current_box[3]), max(self.current_box[0], self.current_box[2]), max(self.current_box[1], self.current_box[3])]# 检查框的有效性box_width = self.current_box[2] - self.current_box[0]box_height = self.current_box[3] - self.current_box[1]print(f"添加的框: {self.current_box}, 尺寸: {box_width}x{box_height}")if box_width >= 5 and box_height >= 5:self.boxes.append(self.current_box.copy())print(f"当前框数量: {len(self.boxes)}")self.current_box = []self.need_update = Truedef update_display(self, frame):"""更新显示图像"""display_frame = frame.copy()if not self.monitoring_mode:# 绘制所有已保存的框for i, box in enumerate(self.boxes):color = (0, 255, 0) if i == len(self.boxes) - 1 else (0, 200, 200)cv2.rectangle(display_frame, (box[0], box[1]), (box[2], box[3]), color, 2)cv2.putText(display_frame, str(i+1), (box[0]+5, box[1]+20), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2)# 绘制当前正在绘制的框if len(self.current_box) == 4:cv2.rectangle(display_frame, (self.current_box[0], self.current_box[1]), (self.current_box[2], self.current_box[3]), (0, 0, 255), 2)cv2.imshow("Camera Feed", display_frame)self.need_update = Falsedef perform_segmentation(self, frame):if not self.boxes:print("没有选择任何框，请先选择框")return None, Noneprint(f"处理 {len(self.boxes)} 个框的分割...")start_time = time.time()try:frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)predictor.set_image(frame_rgb)# 为每个框执行分割all_masks = []all_scores = []for i, box in enumerate(self.boxes):print(f"处理框 {i+1}: {box}")masks, scores, _ = predictor.predict(point_coords=None,point_labels=None,box=np.array(box),multimask_output=False)if len(masks) > 0:all_masks.append(masks[0])all_scores.append(scores[0])if not all_masks:print("未生成任何掩码")return None, Noneend_time = time.time()processing_time = end_time - start_timeprint(f"分割完成，耗时: {processing_time:.2f} 秒")print(f"生成掩码数量: {len(all_masks)}")# 创建纯mask显示图像mask_display = np.zeros_like(frame)# 创建结果图像 - 初始化为原图result_image = frame.copy()for i, (mask, score) in enumerate(zip(all_masks, all_scores)):print(f"处理掩码 {i+1}, 分数: {score:.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.5  # 50%透明度border_width = 2   # 边界宽度# 1. 更新纯mask显示for c in range(3):mask_display[:, :, c][mask] = color[c]# 2. 更新结果图像 - 只在mask区域进行叠加# 创建彩色掩码colored_mask = np.zeros_like(result_image)for c in range(3):colored_mask[:, :, c][mask] = color[c]# 创建alpha通道 - mask区域为0.5，其他区域为0alpha = np.zeros((result_image.shape[0], result_image.shape[1]), dtype=np.float32)alpha[mask] = mask_alphaalpha = np.dstack([alpha]*3)  # 转换为3通道# 只在mask区域进行混合result_image = (result_image * (1 - alpha) + colored_mask * alpha).astype(np.uint8)# 添加边界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"{i+1}:{score:.3f}", (int(center_x), int(center_y)), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)return mask_display, result_imageexcept Exception as e:print(f"分割错误: {e}")import tracebacktraceback.print_exc()return None, Nonedef run(self):print("\n=== 实时摄像头交互式分割程序 ===")print("操作说明:")print("1. 在'Camera Feed'窗口上按住鼠标左键并拖动选择框")print("2. 可以连续选择多个框")print("3. 按 's' 键开始实时监测")print("4. 按 'c' 键清除所有框并返回选择模式")print("5. 按 'r' 键重置所有内容")print("6. 按 ESC 或 'q' 键退出")print("=" + "="*30)try:while True:ret, frame = self.cap.read()if not ret:print("无法从摄像头读取帧")breakself.current_frame = frameif self.need_update:self.update_display(frame)# 如果在监测模式下，定期处理if self.monitoring_mode and (time.time() - self.last_process_time > self.process_interval):mask_display, result_image = self.perform_segmentation(frame)if mask_display is not None:cv2.imshow(self.mask_window_name, mask_display)cv2.imshow(self.result_window_name, result_image)self.result_image = result_imageself.last_process_time = time.time()key = cv2.waitKey(1) & 0xFFif key == 27 or key == ord('q'):  # ESC或q键print("用户请求退出")breakelif key == ord('s'):  # 开始实时监测if not self.boxes:print("请先选择至少一个框")continueprint("进入实时监测模式...")self.monitoring_mode = Trueself.last_process_time = time.time()# 立即处理一次mask_display, result_image = self.perform_segmentation(frame)if mask_display is not None:cv2.imshow(self.mask_window_name, mask_display)cv2.imshow(self.result_window_name, result_image)self.result_image = result_imageelif key == ord('c'):  # 清除所有框并返回选择模式print("清除所有框并返回选择模式")self.boxes = []self.current_box = []self.monitoring_mode = Falseself.need_update = Truecv2.imshow(self.mask_window_name, np.zeros_like(frame))cv2.imshow(self.result_window_name, frame.copy())elif key == ord('r'):  # 重置所有内容print("重置所有内容")self.boxes = []self.current_box = []self.monitoring_mode = Falseself.need_update = Truecv2.imshow("Camera Feed", frame)cv2.imshow(self.mask_window_name, np.zeros_like(frame))cv2.imshow(self.result_window_name, frame.copy())except KeyboardInterrupt:print("程序被中断")finally:self.cap.release()cv2.destroyAllWindows()print("程序结束")# 创建选择器实例并运行
try:selector = CameraBoxSelector()selector.run()
except Exception as e:print(f"初始化错误: {e}")sys.exit(1)# 清理内存
if device.type == "cuda":torch.cuda.empty_cache()