Update app.py

app.py

@@ -281,19 +281,44 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         
         trajectory_img = trajectory_img.cpu().numpy().astype(np.uint8)
     
+    # 修改热力图生成部分
     if torch.cuda.is_available():
-        heatmap = torch.zeros((height, width), device=device)
-        for x, y in filtered_points:
-            # 确保坐标是整数
-            x, y = int(x), int(y)
+        logger.info("使用GPU生成热力图")
+        try:
+            heatmap = torch.zeros((height, width), device=device)
+            for pos in filtered_points:
+                # 确保坐标是整数并且在有效范围内
+                x, y = map(int, pos)  # 明确转换为整数
+                if 0 <= x < width and 0 <= y < height:
+                    temp_heatmap = torch.zeros((height, width), device=device)
+                    temp_heatmap[int(y), int(x)] = 1  # 再次确保是整数
+                    # 使用GPU的高斯模糊
+                    temp_heatmap = gaussian_blur_gpu(temp_heatmap, kernel_size=31, sigma=10)
+                    heatmap += temp_heatmap
+            
+            heatmap = heatmap.cpu().numpy()
+        except Exception as e:
+            logger.error(f"GPU热力图生成失败: {str(e)}")
+            # 回退到CPU处理
+            logger.info("切换到CPU生成热力图")
+            heatmap = np.zeros((height, width), dtype=np.float32)
+            for pos in filtered_points:
+                x, y = map(int, pos)
+                if 0 <= x < width and 0 <= y < height:
+                    temp_heatmap = np.zeros((height, width), dtype=np.float32)
+                    temp_heatmap[int(y), int(x)] = 1
+                    temp_heatmap = cv2.GaussianBlur(temp_heatmap, (31, 31), 10)
+                    heatmap += temp_heatmap
+    else:
+        logger.info("使用CPU生成热力图")
+        heatmap = np.zeros((height, width), dtype=np.float32)
+        for pos in filtered_points:
+            x, y = map(int, pos)
             if 0 <= x < width and 0 <= y < height:
-                temp_heatmap = torch.zeros((height, width), device=device)
-                temp_heatmap[y, x] = 1
-                # 使用GPU的高斯模糊
-                temp_heatmap = gaussian_blur_gpu(temp_heatmap, kernel_size=31, sigma=10)
+                temp_heatmap = np.zeros((height, width), dtype=np.float32)
+                temp_heatmap[int(y), int(x)] = 1
+                temp_heatmap = cv2.GaussianBlur(temp_heatmap, (31, 31), 10)
                 heatmap += temp_heatmap
-        
-        heatmap = heatmap.cpu().numpy()
     
     trajectory_frames = []
     heatmap_frames = []
@@ -372,7 +397,7 @@ def gaussian_blur_gpu(tensor, kernel_size=31, sigma=10):
 def get_gaussian_kernel2d(kernel_size, sigma):
     """生成2D高斯核"""
     kernel_x = torch.linspace(-kernel_size//2, kernel_size//2, kernel_size)
-    x, y = torch.meshgrid(kernel_x, kernel_x)
+    x, y = torch.meshgrid(kernel_x, kernel_x, indexing='ij')
     kernel = torch.exp(-(x.pow(2) + y.pow(2)) / (2 * sigma ** 2))
     return kernel / kernel.sum()
 
@@ -432,7 +457,7 @@ with gr.Blocks() as demo:
         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")
+        gr.Markdown("上传视频来检测和分析小鼠行为 | Upload a video to detect and analyze mice behavior")
         
         with gr.Row():
             with gr.Column():
@@ -473,7 +498,7 @@ with gr.Blocks() as demo:
         ### 使用说明
         1. 上传视频文件
         2. 设置处理参数：
-           - 处理时长：需要分析的视频时长（秒）
+           - 处理时间：需要分析的视频时长（秒）
            - 置信度阈值：检测的置信度要求（越高越严格）
            - 最大检测数量：每帧最多检测的目标数量
         3. 等待处理完成