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,43 +282,66 @@ 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():
-        logger.info("使用GPU生成热力图")
+    def generate_heatmap(points, width, height, device='cuda'):
+        """生成热力图的独立函数"""
         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()
+            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
+                
+                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
+                
+                return heatmap
         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
+            logger.error(f"热力图生成失败: {str(e)}")
+            # 返回空热力图
+            return np.zeros((height, width), dtype=np.float32)
+
+    # 生成热力图
+    logger.info("开始生成热力图")
+    heatmap = generate_heatmap(filtered_points, width, height, device=device)
+
+    # 归一化和着色热力图
+    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:
-        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
+        # 如果热力图全为0，创建一个空白的彩色图像
+        heatmap_colored = np.full((height, width, 3), 255, dtype=np.uint8)
+
+    logger.info("热力图生成完成")
     
     trajectory_frames = []
     heatmap_frames = []