Update app.py

app.py

@@ -10,6 +10,7 @@ import cv2
 from pathlib import Path
 import tempfile
 import imageio
+from tqdm import tqdm  # 新增: 导入tqdm
 
 # 从环境变量获取密码
 APP_USERNAME = "admin"  # 用户名保持固定
@@ -71,7 +72,7 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
     base_size = min(width, height)
     line_thickness = max(1, int(base_size * 0.002))  # 0.2% 的最小边长
 
-    # 设置推理参数并处理视频
+    # 修改: 设置推理参数并处理视频,关闭verbose输出
     results = model.predict(
         source=video_path,
         device=device,
@@ -79,24 +80,25 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         save=False,
         show=False,
         stream=True,
-        line_width=line_thickness,  # 线宽
-        boxes=True,  # 显示边界框
+        line_width=line_thickness,
+        boxes=True,
         show_labels=True,
         show_conf=True,
         vid_stride=1,
         max_det=max_det,
-        retina_masks=True,  # 更精细的显示
-        verbose=False
+        retina_masks=True,
+        verbose=False  # 关闭YOLO默认日志输出
     )
     
     # 处理结果
     frame_count = 0
     detection_info = []
-    
-    # 用于记录轨迹和热图数据
     all_positions = []
     heatmap = np.zeros((height, width), dtype=np.float32)
     
+    # 新增: 创建进度条
+    pbar = tqdm(total=total_frames, desc="处理视频", unit="帧")
+    
     for r in results:
         frame = r.plot()
         
@@ -106,16 +108,13 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
             if isinstance(kpts, torch.Tensor):
                 kpts = kpts.cpu().numpy()
             
-            # 取第一个检测目标的第一个关键点（通常是头部）
             if kpts.shape == (1, 8, 3):  # [num_objects, num_keypoints, xyz]
-                x, y = int(kpts[0, 0, 0]), int(kpts[0, 0, 1])  # 使用第一个关键点
+                x, y = int(kpts[0, 0, 0]), int(kpts[0, 0, 1])
                 all_positions.append([x, y])
                 
-                # 更新热图，使用高斯核来平滑
                 if 0 <= x < width and 0 <= y < height:
-                    # 创建高斯核心点
-                    sigma = 10  # 调整这个值来改变热点大小
-                    kernel_size = 31  # 必须是奇数
+                    sigma = 10
+                    kernel_size = 31
                     temp_heatmap = np.zeros((height, width), dtype=np.float32)
                     temp_heatmap[y, x] = 1
                     temp_heatmap = cv2.GaussianBlur(temp_heatmap, (kernel_size, kernel_size), sigma)
@@ -138,17 +137,17 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
             })
         
         detection_info.append(frame_info)
-        
-        # 写入视频
         video_writer.write(frame)
         
         frame_count += 1
+        pbar.update(1)  # 更新进度条
+        
         if process_seconds and frame_count >= total_frames:
             break
     
-    # 释放视频写入器
+    pbar.close()  # 关闭进度条
     video_writer.release()
-    
+
     # 生成分析报告
     confidences = [float(det['confidence'].strip('%'))/100 for info in detection_info for det in info['detections']]
     hist, bins = np.histogram(confidences, bins=5)
@@ -174,7 +173,6 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
 {confidence_report}
 """
     
-    # 在生成轨迹图之前，添加异常点过滤和轨迹平滑
     def filter_trajectories(positions, width, height, max_jump_distance=100):
         """
         过滤轨迹中的异常点
@@ -193,23 +191,18 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         for i, pos in enumerate(positions):
             x, y = pos
             
-            # 检查点是否在有效范围内
             if not (0 <= x < width and 0 <= y < height):
                 continue
             
-            # 第一个有效点
             if last_valid_pos is None:
                 filtered_positions.append(pos)
                 last_valid_pos = pos
                 continue
             
-            # 计算与上一个点的距离
             distance = np.sqrt((x - last_valid_pos[0])**2 + (y - last_valid_pos[1])**2)
             
             if distance > max_jump_distance:
-                # 如果距离太大，进行插值
                 if len(filtered_positions) > 0:
-                    # 寻找下一个有效点
                     next_valid_pos = None
                     for next_pos in positions[i:]:
                         nx, ny = next_pos
@@ -220,7 +213,6 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
                                 break
                     
                     if next_valid_pos is not None:
-                        # 线性插值
                         steps = max(2, int(distance / max_jump_distance))
                         for j in range(1, steps):
                             alpha = j / steps
@@ -233,22 +225,18 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
                 filtered_positions.append(pos)
                 last_valid_pos = pos
         
-        # 使用移动平均平滑轨迹
         window_size = 5
         smoothed_positions = []
         
         if len(filtered_positions) >= window_size:
-            # 添加开始的点
             smoothed_positions.extend(filtered_positions[:window_size//2])
             
-            # 平滑中间的点
             for i in range(window_size//2, len(filtered_positions) - window_size//2):
                 window = filtered_positions[i-window_size//2:i+window_size//2+1]
                 smoothed_x = int(np.mean([p[0] for p in window]))
                 smoothed_y = int(np.mean([p[1] for p in window]))
                 smoothed_positions.append([smoothed_x, smoothed_y])
             
-            # 添加结束的点
             smoothed_positions.extend(filtered_positions[-window_size//2:])
         else:
             smoothed_positions = filtered_positions
@@ -256,17 +244,14 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         return smoothed_positions
 
     # 修改轨迹图生成部分
-    trajectory_img = np.zeros((height, width, 3), dtype=np.uint8) + 255  # 白色背景
+    trajectory_img = np.zeros((height, width, 3), dtype=np.uint8) + 255
     points = np.array(all_positions, dtype=np.int32)
     if len(points) > 1:
-        # 过滤和平滑轨迹
         filtered_points = filter_trajectories(points.tolist(), width, height)
         points = np.array(filtered_points, dtype=np.int32)
         
-        # 绘制轨迹线，使用渐变色
         for i in range(len(points) - 1):
             ratio = i / (len(points) - 1)
-            # 使用从蓝到红的渐变色
             color = (
                 int((1 - ratio) * 255),  # B
                 50,                      # G
@@ -274,60 +259,48 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
             )
             cv2.line(trajectory_img, tuple(points[i]), tuple(points[i + 1]), color, 2)
         
-        # 绘制起点和终点
-        cv2.circle(trajectory_img, tuple(points[0]), 8, (0, 255, 0), -1)  # 绿色起点
-        cv2.circle(trajectory_img, tuple(points[-1]), 8, (0, 0, 255), -1)  # 红色终点
+        cv2.circle(trajectory_img, tuple(points[0]), 8, (0, 255, 0), -1)
+        cv2.circle(trajectory_img, tuple(points[-1]), 8, (0, 0, 255), -1)
         
-        # 每隔一定帧数添加方向箭头
-        arrow_interval = max(len(points) // 20, 1)  # 控制箭头数量
+        arrow_interval = max(len(points) // 20, 1)
         for i in range(0, len(points) - arrow_interval, arrow_interval):
             pt1 = tuple(points[i])
             pt2 = tuple(points[i + arrow_interval])
-            # 计算箭头方向
             angle = np.arctan2(pt2[1] - pt1[1], pt2[0] - pt1[0])
-            # 绘制箭头
             cv2.arrowedLine(trajectory_img, pt1, pt2, (100, 100, 100), 1, tipLength=0.2)
     
-    # 处理热图
-    if np.max(heatmap) > 0:  # 确保有数据
+    if np.max(heatmap) > 0:
         heatmap_normalized = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX)
         heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), cv2.COLORMAP_JET)
-        # 添加一些透明度
         alpha = 0.7
         heatmap_colored = cv2.addWeighted(heatmap_colored, alpha, np.full_like(heatmap_colored, 255), 1-alpha, 0)
     
-    # 准备GIF动画帧
     trajectory_frames = []
     heatmap_frames = []
     
-    # 创建基础图像
     base_trajectory = np.zeros((height, width, 3), dtype=np.uint8) + 255
     base_heatmap = np.zeros((height, width), dtype=np.float32)
     
-    # 每N帧生成一个动画帧
-    frame_interval = max(1, len(filtered_points) // 50)  # 控制GIF帧数
+    frame_interval = max(1, len(filtered_points) // 50)
     
     for i in range(0, len(filtered_points), frame_interval):
         current_points = filtered_points[:i+1]
         
-        # 轨迹动画帧 - 修改颜色方案，与静态图保持一致
         frame_trajectory = base_trajectory.copy()
         if len(current_points) > 1:
             points = np.array(current_points, dtype=np.int32)
             for j in range(len(points) - 1):
-                ratio = j / (len(current_points) - 1)  # 修改这里，使用当前总长度
+                ratio = j / (len(current_points) - 1)
                 color = (
-                    int((1 - ratio) * 255),  # B
-                    50,                       # G
-                    int(ratio * 255)          # R
+                    int((1 - ratio) * 255),
+                    50,
+                    int(ratio * 255)
                 )
                 cv2.line(frame_trajectory, tuple(points[j]), tuple(points[j + 1]), color, 2)
             
-            # 绘制当前位置点
             cv2.circle(frame_trajectory, tuple(points[-1]), 8, (0, 0, 255), -1)
         trajectory_frames.append(frame_trajectory)
         
-        # 热图动画帧 - 修改颜色方案，使用与静态图相同的配色
         frame_heatmap = base_heatmap.copy()
         for x, y in current_points:
             if 0 <= x < width and 0 <= y < height:
@@ -339,18 +312,15 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         if np.max(frame_heatmap) > 0:
             frame_heatmap_norm = cv2.normalize(frame_heatmap, None, 0, 255, cv2.NORM_MINMAX)
             frame_heatmap_color = cv2.applyColorMap(frame_heatmap_norm.astype(np.uint8), cv2.COLORMAP_JET)
-            # 添加白色背景，与静态图保持一致
             frame_heatmap_color = cv2.addWeighted(frame_heatmap_color, 0.7, np.full_like(frame_heatmap_color, 255), 0.3, 0)
             heatmap_frames.append(frame_heatmap_color)
 
-    # 保存GIF动画 - 修改这部分
-    trajectory_gif_path = output_path.replace('.mp4', '_trajectory.gif')  # 使用.gif后缀
-    heatmap_gif_path = output_path.replace('.mp4', '_heatmap.gif')  # 使用.gif后缀
+    trajectory_gif_path = output_path.replace('.mp4', '_trajectory.gif')
+    heatmap_gif_path = output_path.replace('.mp4', '_heatmap.gif')
     
-    imageio.mimsave(trajectory_gif_path, trajectory_frames, duration=50)  # 50ms per frame
+    imageio.mimsave(trajectory_gif_path, trajectory_frames, duration=50)
     imageio.mimsave(heatmap_gif_path, heatmap_frames, duration=50)
     
-    # 保存图像
     trajectory_path = output_path.replace('.mp4', '_trajectory.png')
     heatmap_path = output_path.replace('.mp4', '_heatmap.png')
     cv2.imwrite(trajectory_path, trajectory_img)
@@ -362,14 +332,12 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
 with gr.Blocks() as demo:
     gr.Markdown("# 🐁 小鼠行为分析 (Mice Behavior Analysis)")
     
-    # 登录界面
     with gr.Group() as login_interface:
         username = gr.Textbox(label="用户名")
         password = gr.Textbox(label="密码", type="password")
         login_button = gr.Button("登录")
         login_msg = gr.Textbox(label="消息", interactive=False)
     
-    # 主界面
     with gr.Group(visible=False) as main_interface:
         gr.Markdown("上传视频来检测和分析小鼠行为 | Upload a video to detect and analyze mice behavior")
         
@@ -412,7 +380,7 @@ with gr.Blocks() as demo:
         ### 使用说明
         1. 上传视频文件
         2. 设置处理参数：
-           - 处理时长：需要分析的视频时长（秒
+           - 处理时长：需要分析的视频时长（秒）
            - 置信度阈值：检测的置信度要求（越高越严格）
            - 最大检测数量：每帧最多检测的目标数量
         3. 等待处理完成
@@ -426,7 +394,6 @@ with gr.Blocks() as demo:
         - 最大检测数量建议根据实际场景设置
         """)
     
-    # 设置事件处理
     login_button.click(
         fn=login,
         inputs=[username, password],
@@ -441,4 +408,4 @@ with gr.Blocks() as demo:
     )
 
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(server_name="0.0.0.0", server_port=7860)