Update app.py

app.py

@@ -50,7 +50,7 @@ def login(username, password):
 @spaces.GPU(duration=300)
 def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8):
     """
-    处理视频并进行鼠检测
+    处理视频并进行小鼠检测
     Args:
         video_path: 输入视频路径
         process_seconds: 处理时长(秒)
@@ -282,63 +282,45 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         trajectory_img = trajectory_img.cpu().numpy().astype(np.uint8)
     
     # 修改热力图生成部分
-    def generate_heatmap(points, width, height, device='cuda'):
-        """生成热力图的独立函数"""
-        try:
-            if device == 'cuda' and torch.cuda.is_available():
-                logger.info("使用GPU生成热力图")
-                heatmap = torch.zeros((height, width), device=device)
-                for point in points:
-                    try:
-                        # 确保坐标是整数
-                        x = int(round(float(point[0])))
-                        y = int(round(float(point[1])))
-                        
-                        if 0 <= x < width and 0 <= y < height:
-                            temp_heatmap = torch.zeros((height, width), device=device)
-                            temp_heatmap[y, x] = 1.0
-                            # 使用GPU的高斯模糊
-                            temp_heatmap = gaussian_blur_gpu(temp_heatmap, kernel_size=31, sigma=10)
-                            heatmap += temp_heatmap
-                    except (ValueError, TypeError, IndexError) as e:
-                        logger.warning(f"跳过无效点 {point}: {str(e)}")
-                        continue
+    def generate_heatmap(points, width, height):
+        """简化版本的热力图生成函数"""
+        logger.info(f"开始生成热力图，共 {len(points)} 个点")
+        
+        # 使用numpy生成热力图
+        heatmap = np.zeros((height, width), dtype=np.float32)
+        
+        for point in points:
+            try:
+                # 确保是整数坐标
+                x = min(max(0, int(point[0])), width-1)
+                y = min(max(0, int(point[1])), height-1)
                 
-                return heatmap.cpu().numpy()
-            else:
-                logger.info("使用CPU生成热力图")
-                heatmap = np.zeros((height, width), dtype=np.float32)
-                for point in points:
-                    try:
-                        x = int(round(float(point[0])))
-                        y = int(round(float(point[1])))
-                        
-                        if 0 <= x < width and 0 <= y < height:
-                            temp_heatmap = np.zeros((height, width), dtype=np.float32)
-                            temp_heatmap[y, x] = 1.0
-                            temp_heatmap = cv2.GaussianBlur(temp_heatmap, (31, 31), 10)
-                            heatmap += temp_heatmap
-                    except (ValueError, TypeError, IndexError) as e:
-                        logger.warning(f"跳过无效点 {point}: {str(e)}")
-                        continue
+                # 在点位置标记1
+                heatmap[y, x] = 1.0
                 
-                return heatmap
-        except Exception as e:
-            logger.error(f"热力图生成失败: {str(e)}")
-            # 返回空热力图
-            return np.zeros((height, width), dtype=np.float32)
+                # 直接使用OpenCV的高斯模糊
+                temp_heatmap = cv2.GaussianBlur(heatmap, (31, 31), 10)
+                heatmap = temp_heatmap
+                
+            except Exception as e:
+                logger.warning(f"处理点 {point} 时出错: {e}")
+                continue
+        
+        logger.info("热力图生成完成")
+        return heatmap
 
     # 生成热力图
     logger.info("开始生成热力图")
-    heatmap = generate_heatmap(filtered_points, width, height, device=device)
+    heatmap = generate_heatmap(filtered_points, width, height)
 
-    # 归一化和着色热力图
+    # 归一化和着色
     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)
+        # 添加白色背景
         heatmap_colored = cv2.addWeighted(heatmap_colored, 0.7, np.full_like(heatmap_colored, 255), 0.3, 0)
     else:
-        # 如果热力图全为0，创建一个空白的彩色图像
+        # 如果没有检测到点，返回白色图像
         heatmap_colored = np.full((height, width, 3), 255, dtype=np.uint8)
 
     logger.info("热力图生成完成")
@@ -503,7 +485,7 @@ with gr.Blocks() as demo:
                     value=1,
                     step=1,
                     label="最大检测数量",
-                    info="每帧最多检测的目标数量"
+                    info="每帧���多检测的目标数量"
                 )
                 process_btn = gr.Button("开始处��")