Create mouse_tracker.py

mouse_tracker.py

@@ -0,0 +1,572 @@
+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}")