Update app.py

app.py

@@ -32,7 +32,7 @@ def login(username, password):
         return gr.update(visible=False), gr.update(visible=True), "登录成功"
     return gr.update(visible=True), gr.update(visible=False), "用户名或密码错误"
 
-@spaces.GPU
+@spaces.GPU(duration=300)
 def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8):
     """
     处理视频并进行小鼠检测
@@ -91,10 +91,27 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
     frame_count = 0
     detection_info = []
     
+    # 用于记录轨迹和热图数据
+    all_positions = []
+    heatmap = np.zeros((height, width), dtype=np.float32)
+    
     for r in results:
-        # 获取绘制了预测结果的帧
         frame = r.plot()
         
+        # 收集位置信息
+        if hasattr(r, 'keypoints') and r.keypoints is not None:
+            for kpts in r.keypoints:
+                if isinstance(kpts, torch.Tensor):
+                    kpts = kpts.cpu().numpy()
+                # 使用第一个关键点（比如头部）作为位置参考
+                if len(kpts) > 0:
+                    pos = kpts[0][:2]  # 取x,y坐标
+                    all_positions.append(pos)
+                    # 更新热图
+                    x, y = int(pos[0]), int(pos[1])
+                    if 0 <= x < width and 0 <= y < height:
+                        heatmap[y, x] += 1
+        
         # 收集检测信息
         frame_info = {
             "frame": frame_count + 1,
@@ -124,21 +141,51 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
     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)
+    
+    confidence_report = "\n".join([
+        f"置信度 {bins[i]:.2f}-{bins[i+1]:.2f}: {hist[i]:3d}个检测 ({hist[i]/len(confidences)*100:.1f}%)"
+        for i in range(len(hist))
+    ])
+    
     report = f"""视频分析报告:
 参数设置:
 - 置信度阈值: {conf_threshold:.2f}
 - 最大检测数量: {max_det}
 - 处理时长: {process_seconds}秒
+
 分析结果:
 - 处理帧数: {frame_count}
 - 平均每帧检测到的老鼠数: {np.mean([info['count'] for info in detection_info]):.1f}
 - 最大检测数: {max([info['count'] for info in detection_info])}
 - 最小检测数: {min([info['count'] for info in detection_info])}
+
 置信度分布:
-{np.histogram([float(det['confidence'].strip('%'))/100 for info in detection_info for det in info['detections']], bins=5)[0].tolist()}
+{confidence_report}
 """
     
-    return output_path, report
+    # 生成轨迹图
+    trajectory_img = np.zeros((height, width, 3), dtype=np.uint8) + 255  # 白色背景
+    points = np.array(all_positions, dtype=np.int32)
+    if len(points) > 1:
+        # 绘制轨迹线
+        cv2.polylines(trajectory_img, [points], False, (0, 0, 255), 2)
+        # 绘制起点和终点
+        cv2.circle(trajectory_img, tuple(points[0]), 5, (0, 255, 0), -1)  # 绿色起点
+        cv2.circle(trajectory_img, tuple(points[-1]), 5, (255, 0, 0), -1)  # 红色终点
+    
+    # 生成热图
+    heatmap_normalized = cv2.normalize(heatmap, None, 0, 255, cv2.NORM_MINMAX)
+    heatmap_colored = cv2.applyColorMap(heatmap_normalized.astype(np.uint8), cv2.COLORMAP_JET)
+    
+    # 保存图像
+    trajectory_path = output_path.replace('.mp4', '_trajectory.png')
+    heatmap_path = output_path.replace('.mp4', '_heatmap.png')
+    cv2.imwrite(trajectory_path, trajectory_img)
+    cv2.imwrite(heatmap_path, heatmap_colored)
+    
+    return output_path, trajectory_path, heatmap_path, report
 
 # 创建 Gradio 界面
 with gr.Blocks() as demo:
@@ -182,6 +229,9 @@ with gr.Blocks() as demo:
             
             with gr.Column():
                 video_output = gr.Video(label="检测结果")
+                with gr.Row():
+                    trajectory_output = gr.Image(label="运动轨迹")
+                    heatmap_output = gr.Image(label="热力图")
                 report_output = gr.Textbox(label="分析报告")
         
         gr.Markdown("""
@@ -212,7 +262,7 @@ with gr.Blocks() as demo:
     process_btn.click(
         fn=process_video,
         inputs=[video_input, process_seconds, conf_threshold, max_det],
-        outputs=[video_output, report_output]
+        outputs=[video_output, trajectory_output, heatmap_output, report_output]
     )
 
 if __name__ == "__main__":