import os import cv2 import numpy as np import matplotlib.pyplot as plt import pandas as pd import collections import tempfile from ultralytics import YOLO import math class MouseTrackerAnalyzer: """基于Ultralytics对象跟踪的鼠强迫游泳实验挣扎度分析器""" def __init__(self, model_path, history_size=5, conf=0.25, iou=0.45, max_det=20, verbose=False): # 初始化模型和参数 self.model = YOLO(model_path, task="segment", verbose=False) self.history_size = history_size self.verbose = verbose # 控制日志输出级别 self.struggle_threshold = 0.3 # 挣扎阈值 # 跟踪相关参数 self.conf = conf # 置信度阈值 self.iou = iou # IOU阈值 self.max_det = max_det # 最大检测数量 # 预设16种固定颜色 (BGR顺序) self.colors = [ (255, 0, 0), # 红 (0, 255, 0), # 绿 (0, 0, 255), # 蓝 (255, 255, 0), # 青 (255, 0, 255), # 洋红 (0, 255, 255), # 黄 (128, 0, 0), # 深红 (128, 0, 128), # 紫 (0, 128, 128), # 青绿 (192, 192, 192),# 银 (128, 128, 128),# 灰 (255, 128, 0), # 橙 (255, 0, 128), # 粉 (0, 128, 255), # 浅蓝 (128, 255, 0), # 黄绿 (0, 255, 128) # 浅绿 ] # 追踪相关 self.prev_masks = {} # 上一帧各 ID 二值掩码 self.histories = {} # 各 ID 分数历史队列 self.track_ids = set() # 所有被跟踪的ID # 视频处理状态 self.cap = None self.writer = None self.frame_id = 0 self.results = [] # 存储每帧结果 self.start_frame = 0 self.end_frame = 0 def init_video(self, video_path, output_path=None, start_frame=0, end_frame=None): """初始化视频处理""" # 打开视频并初始化写出器 self.cap = cv2.VideoCapture(video_path) if not self.cap.isOpened(): raise IOError(f"无法打开视频 {video_path}") # 获取视频属性 width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = self.cap.get(cv2.CAP_PROP_FPS) or 30 self.fps = max(fps, 1.0) # 保存帧率到实例变量，确保至少为1 total_frames = int(self.cap.get(cv2.CAP_PROP_FRAME_COUNT)) if self.verbose: print(f"视频尺寸: {width}x{height}, 帧率: {fps}, 总帧数: {total_frames}") # 设置帧范围 self.start_frame = start_frame self.end_frame = end_frame if end_frame is not None else total_frames - 1 # 确保帧范围有效 if self.start_frame < 0: self.start_frame = 0 if self.end_frame >= total_frames: self.end_frame = total_frames - 1 if self.start_frame > self.end_frame: self.start_frame, self.end_frame = self.end_frame, self.start_frame # 将视频定位到起始帧 if self.start_frame > 0: self.cap.set(cv2.CAP_PROP_POS_FRAMES, self.start_frame) # 如果输出为视频则初始化 VideoWriter if output_path and output_path.lower().endswith(('.mp4', '.avi')): # 使用标准编码器 fourcc = cv2.VideoWriter_fourcc(*'mp4v') # 创建VideoWriter self.writer = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) if self.writer.isOpened(): print(f"成功创建输出视频: {output_path}, 尺寸: {width}x{height}") else: print(f"警告: 无法创建输出视频 {output_path}") # 重置状态 self.frame_id = self.start_frame self.results = [] self.prev_masks.clear() self.histories.clear() self.track_ids.clear() if self.verbose: print(f"视频初始化完成: 总帧数 {total_frames}, 分析范围 {self.start_frame}-{self.end_frame}") return total_frames, self.start_frame, self.end_frame def process_frame(self, frame, frame_id): """处理单帧，返回可视化帧和本帧结果列表""" if self.verbose and frame_id % 10 == 0: print(f"process_frame: 处理帧 {frame_id}") try: # 使用YOLO模型跟踪对象 results = self.model.track( frame, persist=True, # 保持跟踪ID的持久性 conf=self.conf, iou=self.iou, max_det=self.max_det, verbose=False ) # 检查是否有检测结果 frame_results = [] if results[0].boxes is None or len(results[0].boxes) == 0: if self.verbose and frame_id % 50 == 0: print("没有检测到任何对象") return frame.copy(), [] # 处理检测结果 if hasattr(results[0], 'masks') and results[0].masks is not None: # 获取掩码和跟踪ID masks = results[0].masks.data.cpu().numpy() track_ids = results[0].boxes.id if track_ids is None: if self.verbose and frame_id % 50 == 0: print("没有获取到跟踪ID") return frame.copy(), [] track_ids = track_ids.int().cpu().numpy() if self.verbose and frame_id % 50 == 0: print(f"检测到 {len(masks)} 个掩码，{len(track_ids)} 个跟踪ID") # 更新跟踪ID集合 for track_id in track_ids: self.track_ids.add(int(track_id)) # 处理每个跟踪对象 for i, (mask, track_id) in enumerate(zip(masks, track_ids)): track_id = int(track_id) # 二值化掩码 bin_mask = (mask > 0.2).astype(np.uint8) # 应用形态学操作清理掩码 kernel = np.ones((5,5), np.uint8) bin_mask = cv2.morphologyEx(bin_mask, cv2.MORPH_CLOSE, kernel) # 调整掩码尺寸到与原始帧相同 if bin_mask.shape != (frame.shape[0], frame.shape[1]): bin_mask = cv2.resize(bin_mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST) # 计算挣扎度 if track_id in self.prev_masks: prev_mask = self.prev_masks[track_id] # 确保比较的掩码尺寸一致 if prev_mask.shape != bin_mask.shape: prev_mask = cv2.resize(prev_mask, (bin_mask.shape[1], bin_mask.shape[0]), interpolation=cv2.INTER_NEAREST) inter = np.logical_and(prev_mask > 0, bin_mask > 0).sum() union = np.logical_or(prev_mask > 0, bin_mask > 0).sum() iou = inter / union if union > 0 else 0 score = 1 - iou if self.verbose and frame_id % 50 == 0: print(f"跟踪ID {track_id} 挣扎分数: {score:.4f} (IoU: {iou:.4f})") else: score = 0.0 if self.verbose and frame_id % 50 == 0: print(f"跟踪ID {track_id} 初始帧，分数为0") # 保存当前掩码和历史 self.prev_masks[track_id] = bin_mask if track_id not in self.histories: self.histories[track_id] = collections.deque(maxlen=self.history_size) self.histories[track_id].append(score) # 计算挣扎状态 is_struggling = score >= self.struggle_threshold # 计算质心 ys, xs = np.where(bin_mask > 0) if len(xs) > 0: centroid = (int(xs.mean()), int(ys.mean())) else: # 如果掩码为空，使用边界框中心点 box = results[0].boxes[i].xyxy.cpu().numpy()[0] centroid = (int((box[0] + box[2]) / 2), int((box[1] + box[3]) / 2)) # 添加到帧结果 frame_results.append({ 'id': track_id, 'score': float(score), 'centroid': centroid, 'is_struggling': is_struggling }) else: if self.verbose and frame_id % 50 == 0: print("没有检测到任何掩码") return frame.copy(), [] # 可视化 - 在这里创建最终的标注帧 annotated = frame.copy() # 绘制掩码和ID for result in frame_results: track_id = result['id'] color = self.colors[track_id % len(self.colors)] # 绘制掩码 if track_id in self.prev_masks: mask = self.prev_masks[track_id] # 确保掩码与帧大小一致 if mask.shape != (frame.shape[0], frame.shape[1]): mask = cv2.resize(mask, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_NEAREST) mask_overlay = np.zeros_like(frame) mask_overlay[mask > 0] = color # 使用更精确的掩码边缘 contours, _ = cv2.findContours(mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE) cv2.drawContours(annotated, contours, -1, color, 2) # 使用addWeighted进行混合 cv2.addWeighted(annotated, 1.0, mask_overlay, 0.4, 0, annotated) # 在质心位置绘制ID和挣扎状态 centroid = result['centroid'] status_text = "Struggle" if result['is_struggling'] else "Static" cv2.putText(annotated, f"ID:{track_id} {status_text}", (centroid[0], centroid[1]), cv2.FONT_HERSHEY_SIMPLEX, 0.7, color, 2) # 在顶部创建黑色半透明条，显示总结信息 cv2.rectangle(annotated, (0, 0), (frame.shape[1], 40), (0, 0, 0), -1) # 计算挣扎中的老鼠数量 struggling_count = sum(1 for r in frame_results if r['is_struggling']) total_count = len(frame_results) # 显示统计信息 cv2.putText(annotated, f"Total: {total_count} Struggling: {struggling_count}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2) # 最后，由于OpenCV以BGR格式工作，但可能需要RGB格式， # 确保返回的图像是BGR格式（视频写入用BGR，显示用RGB） if annotated.dtype != np.uint8: annotated = annotated.astype(np.uint8) return annotated, frame_results except Exception as e: import traceback if self.verbose: print(f"处理帧时出错: {str(e)}") traceback.print_exc() # 返回原始帧和空结果 return frame.copy(), [] def process_video(self, video_path, output_path=None, start_frame=0, end_frame=None, callback=None): """处理整段视频，可选的回调函数用于更新进度""" # 初始化视频 total_frames, start, end = self.init_video(video_path, output_path, start_frame, end_frame) self.results = [] # 确保结果列表被清空 frame_id = start processed_frames = 0 frames_to_process = end - start + 1 last_progress = -1 # 临时保存一帧，用于调试 debug_frame_saved = False while frame_id <= end: ret, frame = self.cap.read() if not ret: break # 处理当前帧 annotated, frame_res = self.process_frame(frame, frame_id) self.results.append(frame_res) # 将当前帧结果存入results列表 # 保存第一帧用于调试 if not debug_frame_saved and len(frame_res) > 0: debug_frame_path = os.path.join(os.path.dirname(output_path), "debug_frame.jpg") cv2.imwrite(debug_frame_path, annotated) print(f"调试: 保存了标注帧到 {debug_frame_path}") debug_frame_saved = True # 写入输出视频 if self.writer: # 确保帧是BGR格式 if len(annotated.shape) == 3 and annotated.shape[2] == 3: # 如果需要，将RGB转换回BGR (OpenCV使用BGR) # 默认应该已经是BGR，但为了确保 if frame_id == start: print(f"调试: 写入标注帧到视频，形状: {annotated.shape}") try: self.writer.write(annotated) except Exception as e: print(f"调试: 写入帧到视频时出错: {str(e)}") import traceback traceback.print_exc() # 更新进度和回调 processed_frames += 1 progress = int(100 * processed_frames / frames_to_process) if progress != last_progress and callback: callback(progress, annotated, frame_res) last_progress = progress frame_id += 1 # 释放资源 self.cap.release() if self.writer: self.writer.release() print(f"调试: 视频写入完成，保存到: {output_path}") return self.results def save_results(self, csv_path): """导出分析结果到 CSV""" import csv with open(csv_path, 'w', newline='') as f: writer = csv.writer(f) writer.writerow(['frame_id', 'mouse_id', 'score', 'is_struggling']) for fid, frs in enumerate(self.results): for fr in frs: writer.writerow([ fid + self.start_frame, fr['id'], f"{fr['score']:.4f}", 1 if fr.get('is_struggling', False) else 0 ]) def generate_time_series_plot(self, threshold=None): """生成时序图分析""" try: print(f"Starting to generate time series plot with {len(self.results)} frames of data") if not self.results or len(self.results) < 10: print("Not enough data for time series plot (need at least 10 frames)") return None # 使用传入的阈值或默认阈值 if threshold is None: threshold = self.struggle_threshold # 使用保存的帧率，确保不会出现除以零的情况 fps = getattr(self, 'fps', None) if fps is None or fps <= 0: fps = 30 # 使用默认帧率 print(f"Warning: Invalid frame rate detected, using default: {fps} fps") else: print(f"Using frame rate: {fps} fps") # 处理数据 frames = [] mouse_data = {} mouse_positions = {} # 用于存储每只老鼠的平均X坐标 for frame_id, frame_results in enumerate(self.results): frames.append(frame_id + self.start_frame) # 使用真实帧号 for result in frame_results: mouse_id = result['id'] if mouse_id not in mouse_data: mouse_data[mouse_id] = {'frames': [], 'seconds': [], 'scores': [], 'struggling': []} mouse_positions[mouse_id] = [] # 初始化X坐标列表 frame_num = frame_id + self.start_frame second = frame_num / fps # 转换为秒 mouse_data[mouse_id]['frames'].append(frame_num) mouse_data[mouse_id]['seconds'].append(second) mouse_data[mouse_id]['scores'].append(result['score']) mouse_data[mouse_id]['struggling'].append(1 if result.get('is_struggling', False) else 0) # 记录质心的X坐标 if 'centroid' in result: mouse_positions[mouse_id].append(result['centroid'][0]) print(f"Processed data for {len(mouse_data)} mice") if not mouse_data: print("No valid mouse data to plot") return None # 计算每只老鼠的平均X坐标并按从左到右排序 avg_positions = {} for mouse_id, positions in mouse_positions.items(): if positions: avg_positions[mouse_id] = sum(positions) / len(positions) else: avg_positions[mouse_id] = float('inf') # 如果没有位置数据，放到最后 # 按从左到右排序老鼠ID sorted_mice = sorted(mouse_data.keys(), key=lambda mid: avg_positions.get(mid, float('inf'))) print(f"Mice sorted from left to right: {sorted_mice}") # 对数据进行平滑处理 def smooth_data(data, window_size=5): """使用移动平均平滑数据""" if len(data) < window_size: return data smoothed = [] for i in range(len(data)): start = max(0, i - window_size // 2) end = min(len(data), i + window_size // 2 + 1) window = data[start:end] smoothed.append(sum(window) / len(window)) return smoothed # 创建子图 num_mice = len(mouse_data) fig, axes = plt.subplots(num_mice, 1, figsize=(12, 4*num_mice), sharex=True) # 如果只有一只鼠，确保axes是列表 if num_mice == 1: axes = [axes] # 绘制每只老鼠的挣扎得分曲线，按从左到右的顺序 for idx, mouse_id in enumerate(sorted_mice): data = mouse_data[mouse_id] ax = axes[idx] # 平滑数据 smoothed_scores = smooth_data(data['scores'], window_size=5) # 绘制曲线 ax.plot(data['seconds'], smoothed_scores, label=f"Smoothed", color='blue', linewidth=2) ax.plot(data['seconds'], data['scores'], label=f"Raw", color='lightblue', alpha=0.5, linewidth=1) # 标记挣扎区域 for i, is_struggling in enumerate(data['struggling']): if is_struggling: ax.axvspan(data['seconds'][i]-0.5/fps, data['seconds'][i]+0.5/fps, alpha=0.1, color='red') # 绘制阈值线 ax.axhline(y=threshold, color='r', linestyle='--', label=f"Threshold ({threshold:.2f})") # 设置图表 ax.set_ylabel('Struggle Score') position_text = f"(Position: Left #{sorted_mice.index(mouse_id)+1})" if mouse_id in avg_positions else "" ax.set_title(f'Mouse {mouse_id} Struggle Score {position_text}') ax.legend(loc='upper right') ax.grid(True) # 设置Y轴范围0-1 ax.set_ylim(-0.05, 1.05) # 设置共享的X轴标签 axes[-1].set_xlabel('Time (seconds)') # 动态调整x轴范围，精确到0.1秒 if frames: start_time = self.start_frame / fps end_time = max(frames) / fps # 扩展一点范围以便更好地显示 axes[-1].set_xlim(start_time, end_time) # 设置次要刻度（细网格线） tick_interval = 0.1 # 保持0.1秒的细网格 minor_ticks = np.arange(start_time, end_time + tick_interval, tick_interval) axes[-1].set_xticks(minor_ticks, minor=True) # 设置主要刻度（标签和粗网格线）- 整秒 major_start = math.ceil(start_time) major_end = math.floor(end_time) major_ticks = np.arange(major_start, major_end + 1, 1.0) # 整秒刻度 axes[-1].set_xticks(major_ticks) axes[-1].set_xticklabels([f"{int(t)}" for t in major_ticks]) # 整数秒标签 # 设置网格 axes[-1].grid(True, which='both') axes[-1].grid(which='minor', alpha=0.2) axes[-1].grid(which='major', alpha=0.5) plt.tight_layout() # 保存图表到临时文件并返回路径 temp_file = tempfile.NamedTemporaryFile(suffix='.png', delete=False) plt.savefig(temp_file.name, dpi=150, bbox_inches='tight') plt.close() print(f"Time series plot saved to: {temp_file.name}") return temp_file.name except Exception as e: import traceback print(f"Error generating time series plot: {str(e)}") traceback.print_exc() return None if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="鼠强迫游泳实验挣扎度分析") parser.add_argument('--video', type=str, required=True, help='输入视频路径') parser.add_argument('--model', type=str, required=True, help='模型文件路径') parser.add_argument('--output', type=str, help='输出视频路径') parser.add_argument('--csv', type=str, help='输出CSV结果路径') parser.add_argument('--conf', type=float, default=0.25, help='置信度阈值') parser.add_argument('--iou', type=float, default=0.45, help='IOU阈值') parser.add_argument('--max-det', type=int, default=20, help='最大检测数量') parser.add_argument('--threshold', type=float, default=0.3, help='挣扎阈值') parser.add_argument('--start', type=int, default=0, help='起始帧') parser.add_argument('--end', type=int, default=None, help='结束帧') parser.add_argument('--verbose', action='store_true', help='详细输出') args = parser.parse_args() # 设置输出路径 if not args.output: video_name = os.path.splitext(os.path.basename(args.video))[0] args.output = os.path.join(os.path.dirname(args.video), f"{video_name}_out.mp4") if not args.csv: video_name = os.path.splitext(os.path.basename(args.video))[0] args.csv = os.path.join(os.path.dirname(args.video), f"{video_name}_results.csv") # 创建分析器并处理 analyzer = MouseTrackerAnalyzer( model_path=args.model, conf=args.conf, iou=args.iou, max_det=args.max_det, verbose=args.verbose ) analyzer.struggle_threshold = args.threshold # 进度回调函数 def progress_callback(progress, frame, results): print(f"处理进度: {progress}%, 检测到 {len(results)} 个对象") # 处理视频 analyzer.process_video( video_path=args.video, output_path=args.output, start_frame=args.start, end_frame=args.end, callback=progress_callback ) # 保存结果 analyzer.save_results(args.csv) # 生成分析图表 plot_path = analyzer.generate_time_series_plot() if plot_path: print(f"挣扎度时序分析图已保存到: {plot_path}") print(f"分析完成，视频已保存到: {args.output}") print(f"结果数据已保存到: {args.csv}")