Update app.py

app.py

@@ -284,6 +284,8 @@ def process_video(video_path, process_seconds=20, conf_threshold=0.2, max_det=8)
     if torch.cuda.is_available():
         heatmap = torch.zeros((height, width), device=device)
         for x, y in filtered_points:
+            # 确保坐标是整数
+            x, y = int(x), int(y)
             if 0 <= x < width and 0 <= y < height:
                 temp_heatmap = torch.zeros((height, width), device=device)
                 temp_heatmap[y, x] = 1
@@ -376,15 +378,18 @@ def get_gaussian_kernel2d(kernel_size, sigma):
 
 def draw_line_gpu(image, pt1, pt2, color, thickness=1):
     """GPU版本的线段绘制"""
-    x1, y1 = pt1
-    x2, y2 = pt2
+    x1, y1 = map(int, pt1)  # 确保是整数
+    x2, y2 = map(int, pt2)  # 确保是整数
     dx = abs(x2 - x1)
     dy = abs(y2 - y1)
     
-    if dx > dy:
-        steps = dx
-    else:
-        steps = dy
+    # 防止除零错误
+    steps = max(dx, dy)
+    if steps == 0:
+        # 如果是同一个点，直接画点
+        if 0 <= x1 < image.shape[1] and 0 <= y1 < image.shape[0]:
+            image[y1, x1] = color
+        return
     
     x_inc = (x2 - x1) / steps
     y_inc = (y2 - y1) / steps
@@ -398,18 +403,17 @@ def draw_line_gpu(image, pt1, pt2, color, thickness=1):
         x += x_inc
         y += y_inc
     
-    points = points.long()
+    points = points.long()  # 转换为整数类型
     valid_points = (points[:, 0] >= 0) & (points[:, 0] < image.shape[1]) & \
                   (points[:, 1] >= 0) & (points[:, 1] < image.shape[0])
     points = points[valid_points]
     
-    # 确保color是正确的数据类型
-    color = color.to(image.dtype)  # 修改这里，确保颜色张量与图像类型匹配
+    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, dtype=torch.long)  # 修改这里，确保是long类型
+                offset_points = points + torch.tensor([dx, dy], device=device, dtype=torch.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]