Bresenham's line algorithm in PyTorch to connect points in a mask

bresenham_torch.py

def draw_line_in_mask(mask, start_point, end_point):
  # Extract x and y coordinates of start and end points
  x0, y0 = start_point[0], start_point[1]
  x1, y1 = end_point[0], end_point[1]

  # Compute differences between start and end points
  dx = abs(x1 - x0)
  dy = abs(y1 - y0)

  # Determine direction of the line
  sx = 1 if x0 < x1 else -1
  sy = 1 if y0 < y1 else -1

  # Compute the initial error
  error = dx - dy

  # Initialize the current position to the start point
  x, y = x0, y0

  # Loop until we reach the end point
  while x != x1 or y != y1:
      # Append the current position to the list of points
      mask[y, x] = 1

      # Compute the error for the next position
      e2 = 2 * error

      # Determine which direction to move
      if e2 > -dy:
          error = error - dy
          x = x + sx
      if e2 < dx:
          error = error + dx
          y = y + sy

  # Append the final position to the list of points
  mask[y1, x1] = 1

  return mask

def get_box_center(box):
    x_c = (box[2] + box[0]) // 2
    y_c = (box[3] + box[1]) // 2
    return x_c, y_c

start_point = get_box_center(from_box)
end_point = get_box_center(to_box)
mask = self.draw_line_in_mask(mask, start_point, end_point) # mask in [H, W] form