-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathbresenham_torch.py
114 lines (100 loc) · 3.62 KB
/
bresenham_torch.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
"""
Modified from https://code.activestate.com/recipes/578112-bresenhams-line-algorithm-in-n-dimensions/
N-D Bresenham line algo
"""
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
def _bresenhamline_nslope(slope):
"""
Normalize slope for Bresenham's line algorithm.
>>> s = np.array([[-2, -2, -2, 0]])
>>> _bresenhamline_nslope(s)
array([[-1., -1., -1., 0.]])
>>> s = np.array([[0, 0, 0, 0]])
>>> _bresenhamline_nslope(s)
array([[ 0., 0., 0., 0.]])
>>> s = np.array([[0, 0, 9, 0]])
>>> _bresenhamline_nslope(s)
array([[ 0., 0., 1., 0.]])
"""
scale = torch.amax(torch.abs(slope), dim=1).reshape(-1, 1)
zeroslope = (scale == 0).all(1)
scale[zeroslope] = torch.ones(1, dtype=torch.long).to(device)
normalizedslope = slope / scale
normalizedslope[zeroslope] = torch.zeros(slope[0].shape).to(device)
return normalizedslope
def _bresenhamlines(start, end, max_iter):
"""
Returns npts lines of length max_iter each. (npts x max_iter x dimension)
>>> s = np.array([[3, 1, 9, 0],[0, 0, 3, 0]])
>>> _bresenhamlines(s, np.zeros(s.shape[1]), max_iter=-1)
array([[[ 3, 1, 8, 0],
[ 2, 1, 7, 0],
[ 2, 1, 6, 0],
[ 2, 1, 5, 0],
[ 1, 0, 4, 0],
[ 1, 0, 3, 0],
[ 1, 0, 2, 0],
[ 0, 0, 1, 0],
[ 0, 0, 0, 0]],
<BLANKLINE>
[[ 0, 0, 2, 0],
[ 0, 0, 1, 0],
[ 0, 0, 0, 0],
[ 0, 0, -1, 0],
[ 0, 0, -2, 0],
[ 0, 0, -3, 0],
[ 0, 0, -4, 0],
[ 0, 0, -5, 0],
[ 0, 0, -6, 0]]])
"""
if max_iter == -1:
max_iter = torch.amax(torch.amax(torch.abs(end - start), dim=1))
npts, dim = start.shape
nslope = _bresenhamline_nslope(end - start)
# steps to iterate on
stepseq = torch.arange(1, max_iter + 1).to(device)
stepmat = stepseq.repeat(dim, 1) #np.tile(stepseq, (dim, 1)).T
stepmat = stepmat.T
# some hacks for broadcasting properly
bline = start[:, None, :] + nslope[:, None, :] * stepmat
# Approximate to nearest int
bline_points = torch.round(bline).to(start.dtype)
return bline_points
def bresenhamline(start, end, max_iter=5):
"""
Returns a list of points from (start, end] by ray tracing a line b/w the
points.
Parameters:
start: An array of start points (number of points x dimension)
end: An end points (1 x dimension)
or An array of end point corresponding to each start point
(number of points x dimension)
max_iter: Max points to traverse. if -1, maximum number of required
points are traversed
Returns:
linevox (n x dimension) A cumulative array of all points traversed by
all the lines so far.
>>> s = np.array([[3, 1, 9, 0],[0, 0, 3, 0]])
>>> bresenhamline(s, np.zeros(s.shape[1]), max_iter=-1)
array([[ 3, 1, 8, 0],
[ 2, 1, 7, 0],
[ 2, 1, 6, 0],
[ 2, 1, 5, 0],
[ 1, 0, 4, 0],
[ 1, 0, 3, 0],
[ 1, 0, 2, 0],
[ 0, 0, 1, 0],
[ 0, 0, 0, 0],
[ 0, 0, 2, 0],
[ 0, 0, 1, 0],
[ 0, 0, 0, 0],
[ 0, 0, -1, 0],
[ 0, 0, -2, 0],
[ 0, 0, -3, 0],
[ 0, 0, -4, 0],
[ 0, 0, -5, 0],
[ 0, 0, -6, 0]])
"""
# Return the points as a single array
return _bresenhamlines(start, end, max_iter).reshape(-1, start.shape[-1])