Update app.py

app.py

@@ -261,7 +261,7 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
         return filtered_points
 
     # 修改轨迹图生成部分
-    trajectory_img = torch.ones((height, width, 3), device=device) * 255
+    trajectory_img = torch.ones((height, width, 3), device=device, dtype=torch.float32)
     points = np.array(all_positions, dtype=np.int32)
     if len(points) > 1:
         filtered_points = filter_trajectories_gpu(points.tolist(), width, height)
@@ -273,7 +273,7 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
                 int((1 - ratio) * 255),  # B
                 50,                      # G
                 int(ratio * 255)         # R
-            ], device=device)
+            ], device=device, dtype=torch.float32)
             
             # 使用GPU绘制线段
             pt1, pt2 = points[i], points[i + 1]
@@ -403,10 +403,13 @@ def draw_line_gpu(image, pt1, pt2, color, thickness=1):
                   (points[:, 1] >= 0) & (points[:, 1] < image.shape[0])
     points = points[valid_points]
     
+    # 确保color是正确的数据类型
+    color = color.to(image.dtype)  # 修改这里，确保颜色张量与图像类型匹配
+    
     if thickness > 1:
         for dx in range(-thickness//2, thickness//2 + 1):
             for dy in range(-thickness//2, thickness//2 + 1):
-                offset_points = points + torch.tensor([dx, dy], device=device)
+                offset_points = points + torch.tensor([dx, dy], device=device, dtype=torch.long)  # 修改这里，确保是long类型
                 valid_offset = (offset_points[:, 0] >= 0) & (offset_points[:, 0] < image.shape[1]) & \
                              (offset_points[:, 1] >= 0) & (offset_points[:, 1] < image.shape[0])
                 offset_points = offset_points[valid_offset]
@@ -425,7 +428,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():