class test_vectorization_compiled:
def setUp(self):
self.data = np.array([[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0], [1, 0, 1], [1, 1, 0], [1, 1, 1]])
self.kdtree = KDTree(self.data)
def test_single_query(self):
d, i = self.kdtree.query([0, 0, 0])
assert isinstance(d, float)
assert isinstance(i, int)
def test_vectorized_query(self):
d, i = self.kdtree.query(np.zeros((2, 4, 3)))
assert_equal(np.shape(d), (2, 4))
assert_equal(np.shape(i), (2, 4))
def test_single_query_multiple_neighbors(self):
s = 23
kk = self.kdtree.n + s
d, i = self.kdtree.query([0, 0, 0], k=kk)
assert_equal(np.shape(d), (kk,))
assert_equal(np.shape(i), (kk,))
assert np.all(~np.isfinite(d[-s:]))
assert np.all(i[-s:] == self.kdtree.n)
def test_vectorized_query_multiple_neighbors(self):
s = 23
kk = self.kdtree.n + s
d, i = self.kdtree.query(np.zeros((2, 4, 3)), k=kk)
assert_equal(np.shape(d), (2, 4, kk))
assert_equal(np.shape(i), (2, 4, kk))
assert np.all(~np.isfinite(d[:, :, -s:]))
assert np.all(i[:, :, -s:] == self.kdtree.n)
class test_small(ConsistencyTests):
def setUp(self):
self.data = np.array([[0,0,0],
[0,0,1],
[0,1,0],
[0,1,1],
[1,0,0],
[1,0,1],
[1,1,0],
[1,1,1]])
self.kdtree = KDTree(self.data)
self.n = self.kdtree.n
self.m = self.kdtree.m
np.random.seed(1234)
self.x = np.random.randn(3)
self.d = 0.5
self.k = 4
def test_nearest(self):
assert_array_equal(
self.kdtree.query((0,0,0.1), 1),
(0.1,0))
def test_nearest_two(self):
assert_array_equal(
self.kdtree.query((0,0,0.1), 2),
([0.1,0.9],[0,1]))
class FeaturesQuery(object):
def __init__(self, fgroup, coord_modes=['lm']):
self.features1 = fgroup.features
self.nf1 = len(self.features1)
coords1 = np.array([f.get_coord(mode=mode) for mode in coord_modes for f in self.features1])
if len(self.features1) > 0:
self.kdtree = KDTree(coords1)
def get_features(self):
return set(self.features1)
def find(self, feature2, tol=np.inf):
if len(self.features1) == 0:
return []
d, i = self.kdtree.query(feature2.get_coord(), k=len(self.features1),
distance_upper_bound=tol)
if len(self.features1) == 1:
if d < tol:
return self.features1
return []
i = i[d < tol]
return nputils.uniq([self.features1[k % self.nf1] for k in i])
def get_match(self, fgroup2, tol=2):
if len(self.features1) == 0:
return FeaturesMatch()
coords2 = np.array([f.get_coord() for f in fgroup2.features])
d, i = self.kdtree.query(coords2, k=1, distance_upper_bound=tol)
i1 = i[d < tol]
i2 = np.arange(len(coords2))[d < tol]
features1 = [self.features1[k % self.nf1] for k in i1]
features2 = [fgroup2.features[k] for k in i2]
return FeaturesMatch(features1=features1, features2=features2)
def calc_nearest_neighbor_distances(self, detail_level):
rf = self.get_detail_level_results(detail_level)
if rf.is_event_list_column_name('eNN_dist') and \
rf.is_event_list_column_name('eNN_dist') and \
rf.is_event_list_column_name('eNN_dist'):
return rf
cons_points = numpy.array([e.position() for e in rf.consumers()])
prod_points = numpy.array([e.position() for e in rf.producers()])
if self.wrapping:
cons_points = wrap_points(cons_points, self.side_length, is_3d=self.is_3d)
prod_points = wrap_points(prod_points, self.side_length, is_3d=self.is_3d)
cons_tree = KDTree(cons_points)
prod_tree = KDTree(prod_points)
for e in rf.events:
c_dist = cons_tree.query(e.position(), k=2)[0]
p_dist= prod_tree.query(e.position(), k=2)[0]
if (e.vars['mark'] == 2):
e.vars['sNN_dist'] = c_dist[1]
e.vars['oNN_dist'] = p_dist[0]
else:
e.vars['sNN_dist'] = p_dist[1]
e.vars['oNN_dist'] = c_dist[0]
e.vars['eNN_dist'] = min(e.vars['sNN_dist'], e.vars['oNN_dist'])
rf.add_event_list_column_name('eNN_dist')
rf.add_event_list_column_name('oNN_dist')
rf.add_event_list_column_name('sNN_dist')
#rf.eventList.update_text()
rf.write()
return rf
def query_tree3(table_j, table_ks, table_h, max_sep=0.0001):
"""
For two pandas database files, with the ra/dec header names x_wcs/y_wcs as defined in read_sex
"""
tree2 = KDTree(zip(table_ks['X_WORLD'].ravel(), table_ks['Y_WORLD'].ravel()))
tree3 = KDTree(zip(table_h['X_WORLD'].ravel(), table_h['Y_WORLD'].ravel()))
idxs1 = []
idxs2 = []
idxs3 = []
dups2 = []
dups3 = []
for i in range(len(table_j)):
d2, icoords2 = tree2.query((table_j['X_WORLD'][i], table_j['Y_WORLD'][i]), distance_upper_bound=max_sep)
d3, icoords3 = tree3.query((table_j['X_WORLD'][i], table_j['Y_WORLD'][i]), distance_upper_bound=max_sep)
if (d2 != np.inf) & (d3 != np.inf):
if icoords2 in idxs2:
dups2.append(icoords2)
if icoords3 in idxs3:
dups3.append(icoords3)
idxs1.append(i)
idxs2.append(icoords2)
idxs3.append(icoords3)
if (len(dups2) > 0) | (len(dups3) > 0):
raise Exception, 'WARNING: duplicates found'
return np.array((idxs1, idxs2, idxs3))
def compute_lfs(datadict, k=10):
from scipy.spatial import KDTree
# collect all ma_coords that are not NaN
ma_coords = np.concatenate([datadict['ma_coords_in'], datadict['ma_coords_out']])
ma_coords = ma_coords[~np.isnan(ma_coords).any(axis=1)]
# build kd-tree of ma_coords to compute lfs
pykdtree = KDTree(ma_coords)
if k > 1:
datadict['lfs'] = np.sqrt(np.median(pykdtree.query(datadict['coords'], k)[0], axis=1))
else:
datadict['lfs'] = np.sqrt(pykdtree.query(datadict['coords'], k)[0])
def plot_lsh_performance_for_bins_projs(n_samples, len_x, k_neighbours, bins_list, projs_list):
d = create_data(n_samples, len_x)
q = create_data(1, len_x)[0]
fig = plt.figure(figsize=(10,10))
ax1 = fig.add_subplot(221)
ax2 = fig.add_subplot(222)
ax3 = fig.add_subplot(223)
ax4 = fig.add_subplot(224)
total,c = len(bins_list)*len(projs_list),1
for p in projs_list:
bOR, bAND, tOR, tAND = [], [], [], []
for b in bins_list:
if c%(total/10)==0:
print "%d"%(c*100./total)+"%",
c+=1
brOR, brAND, trOR, trAND = [], [], [], []
for i in range(10):
q = (np.random.random(size=(len_x))-0.5)*.7
kt = KDTree(d)
ktr = kt.query(q,k_neighbours)
drref = np.max(ktr[0])
lsh = DiscretizedLSH(d,p,b)
ts = time(); lshrAND = lsh.search_AND(q,k_neighbours); teAND = time()-ts
ts = time(); lshrOR = lsh.search_OR(q,k_neighbours); teOR = time()-ts
brAND.append(np.mean([dratio(d[i],q,drref) for i in lshrAND]))
brOR.append(np.mean([dratio(d[i],q,drref) for i in lshrOR]))
trAND.append(teAND)
trOR.append(teOR)
bOR.append(np.mean(brOR))
bAND.append(np.mean(brAND))
tOR.append(np.mean(trOR))
tAND.append(np.mean(trAND))
ax1.plot(bins_list, bOR, label="np "+str(p))
ax2.plot(bins_list, bAND, label="np "+str(p))
ax3.plot(bins_list, tOR, label="np "+str(p))
ax4.plot(bins_list, tAND, label="np "+str(p))
ax1.set_xlabel("bin width")
ax1.set_ylabel("distance ratio")
ax1.set_title("lsh OR (%d,%d) k=%d"%(n_samples, len_x, k_neighbours))
ax1.legend()
ax2.set_xlabel("bin width")
ax2.set_ylabel("distance ratio")
ax2.set_title("lsh AND (%d,%d) k=%d"%(n_samples, len_x, k_neighbours))
ax2.legend()
ax3.set_xlabel("bin width")
ax3.set_ylabel("query time")
ax3.set_title("lsh OR (%d,%d) k=%d"%(n_samples, len_x, k_neighbours))
ax3.legend()
ax4.set_xlabel("bin width")
ax4.set_ylabel("query time")
ax4.set_title("lsh AND (%d,%d) k=%d"%(n_samples, len_x, k_neighbours))
ax4.legend()
def plot_lsh_performance_for_datasizes (row_list, col_list, k_neighbours, bin_width, projs):
fig = plt.figure(figsize=(10,10))
ax1 = fig.add_subplot(221)
ax2 = fig.add_subplot(222)
ax3 = fig.add_subplot(223)
ax4 = fig.add_subplot(224)
total,c = len(row_list)*len(col_list),1
for col in col_list:
bOR, bAND, tOR, tAND = [], [], [], []
for row in row_list:
if c%(total/10)==0:
print "%d"%(c*100./total)+"%",
c+=1
brOR, brAND, trOR, trAND = [], [], [], []
for i in range(10):
d = create_data(row, col)
q = create_data(1, col)[0]
kt = KDTree(d)
ktr = kt.query(q,k_neighbours)
drref = np.max(ktr[0])
lsh = DiscretizedLSH(d,projs,bin_width)
ts = time(); lshrAND = lsh.search_AND(q,k_neighbours); teAND = time()-ts
ts = time(); lshrOR = lsh.search_OR(q,k_neighbours); teOR = time()-ts
brAND.append(np.mean([dratio(d[i],q,drref) for i in lshrAND]))
brOR.append(np.mean([dratio(d[i],q,drref) for i in lshrOR]))
trAND.append(teAND)
trOR.append(teOR)
bOR.append(np.mean(brOR))
bAND.append(np.mean(brAND))
tOR.append(np.mean(trOR))
tAND.append(np.mean(trAND))
ax1.plot(row_list, bOR, label="cols "+str(col))
ax2.plot(row_list, bAND, label="cols "+str(col))
ax3.plot(row_list, tOR, label="cols "+str(col))
ax4.plot(row_list, tAND, label="cols "+str(col))
ax1.set_xlabel("rows")
ax1.set_ylabel("distance ratio")
ax1.set_title("lsh OR bin_w=%.3f, projs=%d, k=%d"%(bin_width, projs, k_neighbours))
ax1.legend()
ax2.set_xlabel("rows")
ax2.set_ylabel("distance ratio")
ax2.set_title("lsh AND bin_w=%.3f, projs=%d, k=%d"%(bin_width, projs, k_neighbours))
ax2.legend()
ax3.set_xlabel("rows")
ax3.set_ylabel("query time")
ax3.set_title("lsh OR bin_w=%.3f, projs=%d, k=%d"%(bin_width, projs, k_neighbours))
ax3.legend()
ax4.set_xlabel("rows")
ax4.set_ylabel("query time")
ax4.set_title("lsh AND bin_w=%.3f, projs=%d, k=%d"%(bin_width, projs, k_neighbours))
ax4.legend()
def treequerycalc(df, k):
points = list(zip(df.X, df.Y)) # make a list of X-Y pairs (lists) from df.X and df.Y
tree = KDTree(points)
[nndist, indices] = tree.query(points, k)
return nndist, indices
def get_query(self, data_path, i, j):
'''
'''
print 'Probing', i, j
#
# create kd tree
#
data = self._kd_tree_data[data_path]
data_without_key = [v[:2] for v in data]
kd_tree = KDTree(data_without_key)
# now look up the closest point to i,j
query_result = kd_tree.query([float(i), float(j)])
print 'Found', query_result[1]
tile = data[query_result[1]]
print 'Found tile', tile
# # now look up this border point in out bbox-to-tile dict
# tile = self._bbox_to_tile[data_path][str(point[0])+'-'+str(point[1])]
# print 'Finding tile', tile
return json.dumps(tile[2])
def warpCloud( xyc, sourceGridPoints, targetGridPoints, warpQuality=9 ):
sourceTree = KDTree(sourceGridPoints, leafsize=10)
warpedXYC = []
for c in xyc:
nearestEdge = sourceTree.query(c,k=warpQuality)
nx = 0.0
ny = 0.0
ws = 0.0
for i in range(warpQuality):
p = targetGridPoints[nearestEdge[1][i]]
w = nearestEdge[0][i]
if w == 0.0:
nx = p[0]
ny = p[1]
ws = 1.0
break
else:
w = 1.0 / w
nx += w * p[0]
ny += w * p[1]
ws += w
warpedXYC.append([nx/ws,ny/ws])
warpedXYC = np.array(warpedXYC)
return warpedXYC
def search_headgroup_loss(filename, x_dim):
from scipy.spatial import KDTree
mz_tol = 0.005
rt_tol = 1
tlc_dataset = TLCdataset(filename, x_dim=x_dim)
tlc_dataset.get_features_from_mean_spectrum(ppm=6.)
feature_array = [[f[0], f[1]] for f in tlc_dataset.feature_list]
kdtree = KDTree(feature_array)
neutral_losses_rt = {}
neutral_losses_i = {}
for head_group in neutral_losses_pos:
neutral_losses_rt[head_group]=[]
neutral_losses_i[head_group] = []
for f in tlc_dataset.feature_list:
for head_group in neutral_losses_pos.keys():
neutral_loss_mz = f[0]-neutral_losses_pos[head_group]
nearest_neighbour = kdtree.query([neutral_loss_mz, f[1]])
if all([np.abs(nearest_neighbour[0]-neutral_loss_mz < mz_tol), np.abs(f[1]-nearest_neighbour[1]<rt_tol)]):
neutral_losses_rt[head_group].append(f[1])
neutral_losses_i[head_group].append(f[2])
plt.figure()
for head_group in neutral_losses_rt:
h = np.histogram(neutral_losses_rt[head_group], bins=tlc_dataset.x_pos, weights=neutral_losses_i[head_group])[0]
plt.plot(h, label=head_group)
plt.legend()
plt.show()
def process_image(image_name, colours, thumbs, thumbsize):
# KD Tree for Nearest-Neighbour
kdt = KDTree(colours)
# open image and make it smaller
img = Image.open(image_name).convert('RGBA')
img.thumbnail((MAX_RES, MAX_RES))
img_data = np.array(img)
# initialise new image
new_shape = (
img_data.shape[0] * thumbsize,
img_data.shape[1] * thumbsize,
img_data.shape[2]
)
new_image = np.zeros(new_shape, dtype=np.float)
# populate new image
for x in range(img_data.shape[0]):
for y in range(img_data.shape[1]):
_, nn_idx = kdt.query(img_data[x, y, :], k=3)
thumb_img = thumbs[random.choice(nn_idx)] # TODO: change to distance-weighted probabilities
x_offset = int((thumbsize - thumb_img.shape[0]) / 2)
y_offset = int((thumbsize - thumb_img.shape[1]) / 2)
x1 = (x * thumbsize) + x_offset
x2 = x1 + thumb_img.shape[0]
y1 = (y * thumbsize) + y_offset
y2 = y1 + thumb_img.shape[1]
new_image[x1:x2, y1:y2, :] = thumb_img
return new_image
def nearest_neighbors(X, k, phi):
"""
Construct pairwise weights by finding the k nearest neighbors to each point
and assigning a Gaussian-based distance.
In particular, w_{i,j} = exp(-phi ||x_i - x_j||_2^2) if j is one of i's k
closest neighbors.
If j is one of i's closest neighbors and i is one of j's closest members, the
edge will only appear once with i < j.
Parameters
----------
X : [n_samples, n_dim] array
k : int
number of neighbors for each sample in X
phi : float
non-negative integer dictating how much weight to assign between pairs of points
"""
from scipy.spatial import KDTree
tree = KDTree(X)
weights = []
for i, x in enumerate(X):
distances, indices = tree.query(x, k+1)
for (d, j) in zip(distances, indices)[1:]:
d = np.exp( -phi * d * d )
if d == 0: continue
if i < j:
weights.append( (i, j, d) )
else:
weights.append( (j, i, d) )
weights = sorted(weights, key=lambda r: (r[0], r[1]))
return unique_rows(np.asarray(weights))
class Palette(object):
def __init__(self, images):
self.cache = {}
self.kdtree = KDTree([v["color"] for v in images.itervalues()])
self.pmap = {}
for (img_path, i) in images.iteritems():
color = tuple(i["color"])
if not color in self.pmap:
self.pmap[color] = deque()
self.pmap[color].append(i["thumbnail"])
def get_closest_color(self, c):
dist, idx = self.kdtree.query(c)
return tuple(self.kdtree.data[idx])
def get_image(self, color):
variants = self.pmap[color]
# Try to use different images
iname = variants.popleft()
variants.append(iname)
img = self.cache.get(iname)
if not img:
img = Image.open(iname)
self.cache[iname] = img
return img
def get_closest_image(self, c):
nc = self.get_closest_color(c)
return self.get_image(nc)
def nearest_neighbors(n,points,qpoints=None,return_distances=False,slow=False):
extra = 0
if qpoints==None:
qpoints=points
if len(points)>1:
extra=1
elif return_distances:
return array([]),array([])
else:
return array([])
#end if
#end if
if n>len(qpoints)-extra:
print 'nearest_neighbors Error: requested more than the total number of neighbors\n maximum is: {0}\n you requested: {1}\nexiting.'.format(len(qpoints)-extra,n)
exit()
#end if
slow = slow or scipy_unavailable
if not slow:
kt = KDTree(points)
dist,ind = kt.query(qpoints,n+extra)
else:
dtable,order = distance_table(points,qpoints,ordering=2)
dist = dtable[:,0:n+extra]
ind = order[:,0:n+extra]
#end if
if extra==0 and n==1 and not slow:
nn = atleast_2d(ind).T
else:
nn = ind[:,extra:]
#end if
if not return_distances:
return nn
else:
return nn,dist
def fit_mesh(head_mesh,scan_mesh,base,tofit_verts,init_coefs = None,constrained = False,regul = None,niter = 1):
if isinstance(tofit_verts,str):
tofit_verts = np.loadtxt(tofit_verts,'int')
if isinstance(base,str):
base = TargetBase(prefix = prefix)
verts = base.vert_list
ntargs = len(base.targets)
m = head_mesh.copy()
nverts = len(scan_mesh)
if init_coefs is not None :
init_target = base.combine_targets(init_coefs)
m[verts]+=init_target
kd = KDTree(scan_mesh)
for iter in xrange(niter):
target = np.zeros((nverts,3))
dist,indx = kd.query(m[tofit_verts])
target[tofit_verts] = scan_mesh[indx] - m[tofit_verts]
target = target[verts]
if constrained :
coefs = base.compute_combinaison_safe(target,rcond,regul = regul)
proj = base.combine_targets(coefs)
else :
proj = base.project_target(target)
if niter > 1 : m[verts]+=(1./niter)*(niter)**(float(iter)/(niter-1)) * proj
else : m[verts]+= proj
return m-head_mesh
def ICP(base, target):
geoCenterTarget = np.mean(target, axis=1)
R = np.eye(3)
t = np.array([0,0,0])
KDindex = KDTree(target.T)
lastDistance = np.array([1,1,1])
sample = base
for i in range(150):
transformedSample = np.dot(sample.T, R) + t
I = KDindex.query(transformedSample)
idx = I[1]
distance = np.sqrt(np.sum(I[0])**2 / transformedSample.shape[0])
if np.sum(np.abs(lastDistance - distance)) < 1e-7:
break
lastDistance = distance
geoCenterTransformedSample = transformedSample.mean(axis=0)
geoCenterTarget = target[:, idx].T.mean(axis = 0)
#Normalize
A = np.dot( (transformedSample - geoCenterTransformedSample).T, (target[:, idx].T - geoCenterTarget) )
U, s, V = np.linalg.svd(A)
R = np.dot(U,V)
geoCentersample = sample.mean(axis=1)
t = geoCenterTarget - np.dot(geoCentersample, R)
return R, t, distance
def generate_capitals(self, map_obj):
kdtree = KDTree([center.point for center in map_obj.centers])
centers_index = {}
for center in map_obj.centers:
centers_index[key(center.point)] = center
height = self.cell_size
width = math.sqrt(3) / 2 * height
row = 0
x = height / 2
while x < 1:
if row % 2 == 0:
y = width
else:
y = width / 2
while y < 1:
_, index = kdtree.query([x, y])
point = kdtree.data[index]
center = centers_index[key(point)]
# Do not put capitals close on to other
is_neighbor_capital = any([neighbor.region for neighbor in center.neighbors])
if not center.water and not is_neighbor_capital:
center.region = Region(center)
map_obj.regions.append(center.region)
y += width
row += 1
x += (height * 3 / 4)
def get_velocity(self, longitude, latitude, height, date):
"""
Find plume velocity that is closest to the given location and time.
:type longitude: float
:param longitude: Longitude of the point of interest.
:type latitude: float
:param latitude: Latitude of the point of interest.
:type altitude: float
:param altitude: Altitude in meters of the point of interest.
:type date: str
:param date: Date of interest formatted according to the ISO8601
standard.
"""
from scipy.spatial import KDTree
from spectroscopy.util import parse_iso_8601
_d = parse_iso_8601(date)
_ts = calendar.timegm((_d.utctimetuple()))
# find nearest point
_t = np.atleast_2d(self.time).T
# TODO: this needs to be changed to account for regular and irregular
# grids
a = np.append(self.position, _t, axis=1)
tree = KDTree(a, leafsize=a.shape[0] + 1)
point = [longitude, latitude, height, _ts]
distances, ndx = tree.query([point], k=1)
vx = self.vx[ndx[0]]
vx_error = self.vx_error[ndx[0]]
vy = self.vy[ndx[0]]
vy_error = self.vy_error[ndx[0]]
vz = self.vz[ndx[0]]
vz_error = self.vz_error[ndx[0]]
time = self.time[ndx[0]]
lon, lat, hght = self.position[ndx[0], :]
return (lon, lat, hght, time, vx, vx_error, vy, vy_error, vz, vz_error)
class Regression(object):
"""
Performs kNN regression
"""
def __init__(self):
self.k = 5
self.metric = np.mean
self.kdtree = None
self.houses = None
self.values = None
def set_data(self, houses, values):
"""
Sets houses and values data
:param houses: pandas.DataFrame with houses parameters
:param values: pandas.Series with houses values
"""
self.houses = houses
self.values = values
self.kdtree = KDTree(self.houses)
def regress(self, query_point):
"""
Calculates predicted value for house with particular parameters
:param query_point: pandas.Series with house parameters
:return: house value
"""
_, indexes = self.kdtree.query(query_point, self.k)
value = self.metric(self.values.iloc[indexes])
if np.isnan(value):
raise Exception('Unexpected result')
else:
return value
class RtreeKDTree(object):
""" wrap scipy KDTree spatial index to look as much like Rtree class
as possible
"""
def __init__(self,stream,interleaved=False):
""" stream: an iterable, returning tuples of the form (id,[xmin,xmax,ymin,ymax],object)
for now, requires that xmin==xmax, and ymin==ymax
For now, only interleaved=False is supported.
"""
if interleaved:
raise Exception,"No support for interleaved index. You must use xxyy ordering"
it = iter(stream)
self.points = []
self.data = []
for fid,xxyy,obj in it:
if xxyy[0]!=xxyy[1] or xxyy[2]!=xxyy[3]:
raise Exception,"No support in kdtree for finite sized objects"
self.points.append([xxyy[0],xxyy[2]])
self.data.append( (fid,obj) )
self.refresh_tree()
def refresh_tree(self):
self.kdt = KDTree(self.points)
def nearest(self, rect, count):
xy = [rect[0],rect[2]]
dists,results = self.kdt.query( xy,k=count )
return [self.data[r][0] for r in results]
def intersects(self,xxyy):
""" This should be made compatible with the regular RTree call...
"""
raise Exception,"Intersects is not implemented in scipy.KDTree"
return []
def insert(self, feat_id, rect=None ):
if rect[0]!=rect[1] or rect[2]!=rect[3]:
raise Exception,"No support in kdtree for finite sized objects"
if rect is None:
raise Exception,"Not sure what inserting an empty rectangle is supposed to do..."
self.points.append( [rect[0],rect[2]] )
self.data.append( [feat_id,None] )
self.refresh_tree()
def feat_id_to_index(self,feat_id):
for i in range(len(self.data)):
if self.data[i][0] == feat_id:
return i
raise Exception,"feature id not found"
def delete(self, feat_id, rect ):
index = self.feat_id_to_index(feat_id)
del self.data[index]
del self.points[index]
self.refresh_tree()
def cluster(self, data): # maps data point to cluster index; optimizes stopping criteria data_length = len(data) cluster_map = [-1] * data_length # start off with a few random cluster centers; # this is the Forgy method and usually OK for standard k-means. current_cluster_centers = data[np.random.choice(data_length, size=self.k, replace=False), :] # repeat the loop until the cluster assignments do not change anymore at_least_one_assignment_changed = True while at_least_one_assignment_changed: # consult a KD tree for performance improvement cluster_center_tree = KDTree(current_cluster_centers, leafsize=5) _, data_membership = cluster_center_tree.query(data) # update stopping condition # manual loop instead of equality operator as the latter always checks all elements, ... at_least_one_assignment_changed = False for i in range(data_length): if cluster_map[i] == data_membership[i]: continue at_least_one_assignment_changed = True break # and finally update the cluster centers to be the centroids of the voronoi subsets cluster_map = data_membership current_cluster_centers = [np.median(data[cluster_map == i, :], axis=0) for i in range(self.k)] return Clusters(data = data, membership = cluster_map, cluster_centers = current_cluster_centers)
def to_mask(self, lons_2d_grid, lats_2d_grid):
"""
:param lons_2d_grid:
:param lats_2d_grid:
:return: the mask of the subregion corresponding to the grid with the upper right and lower left points from self
"""
x_g, y_g, z_g = lat_lon.lon_lat_to_cartesian(lons_2d_grid.flatten(), lats_2d_grid.flatten())
ktree = KDTree(list(zip(x_g, y_g, z_g)))
ll_x, ll_y, ll_z = lat_lon.lon_lat_to_cartesian(self.lleft_lon, self.lleft_lat)
ur_x, ur_y, ur_z = lat_lon.lon_lat_to_cartesian(self.uright_lon, self.uright_lat)
i_g, j_g = np.indices(lons_2d_grid.shape)
i_g_flat, j_g_flat = i_g.flatten(), j_g.flatten()
_, ind_ll = ktree.query((ll_x, ll_y, ll_z), k=1)
_, ind_ur = ktree.query((ur_x, ur_y, ur_z), k=1)
i_ll, j_ll = i_g_flat[ind_ll], j_g_flat[ind_ll]
i_ur, j_ur = i_g_flat[ind_ur], j_g_flat[ind_ur]
res = np.zeros_like(lons_2d_grid, dtype=bool)
res[i_ll:i_ur + 1, j_ll: j_ur + 1] = 1
return res, (i_ll, j_ll), (i_ur, j_ur)
class kNN:
"""A k-Nearest-Neighbours classifier (implemented using KDTrees)"""
def __init__(self, k = 1):
self._k = k # number of nearest neighbours
def train(self, X, y):
# A KDTree is a datastructure that is aware of the spatial
# distribution of its elements and can be used for efficient
# nearest neighbour queries
self._T = KDTree(X)
self._y = y
self._K = np.max(y) + 1 # number of classes
def classify(self, x):
# query the KDTree
dist, ind = self._T.query(x, k=self._k)
# for k == 1, KDTree.query returns only a value, not a list
if self._k == 1:
return self._y[ind]
# count occurences of each class
count = np.zeros(self._K)
for i in ind:
count[self._y[i]] += 1
if np.size(count[count == np.max(count)]) == 1:
# if there is a single class with the most nearest
# neighbours, return it
return np.argmax(count)
else:
# if there are multiple classes with nearest neighbours,
# return the one with smallest mean distance
cls = np.argwhere(count == np.max(count)).flatten()
return cls[np.argmin([np.mean(dist[ind == cl]) for cl in cls])]
def kdtreeApproach(intersections, tripLocations):
#counts is a dictionary that has as keys the intersection index in the intersections list
#and as values the number of trips in the tripLocation list which has the key as the closest
#intersection.
counts = {}
startTime = time.time()
tree = KDTree(intersections)
points = tree.query(tripLocations,k = 1)
indeces = points[1]
counts = Counter(indeces)
#key = indeces.keys()
#value = key[indeces.values().index(max(indeces.values()))]
#print value
#print intersections[value]
#print counts.keys()
#busy = indeces.keys().index(value.index(max(value)))
#print busy
#TODO: insert your code here. You should build the kdtree and use it to query the closest
# intersection for each trip
#
endTime = time.time()
print 'The kdtree computation took', (endTime - startTime), 'seconds'
return counts
class PolygonMap(object):
def __init__(self, seed=None, seed_size=100, relax_iter=4):
"""
"""
if seed is None:
seed = np.random.random((seed_size, 2))
self.points = relax(seed, relax_iter)
self.polymap = Voronoi(self.points)
self.dist_map = KDTree(self.points)
self.poly_graph = nx.Graph()
self.corner_graph = nx.Graph()
for p in range(len(self.polymap.points)):
f = PolyFeature(idx=p, parent=self, center=self.points[p])
self.poly_graph.add_node(p, feature=f)
for e in self.polymap.ridge_points:
self.poly_graph.add_edge(*e.tolist())
for c in range(len(self.polymap.vertices)):
self.corner_graph.add_node(c)
for r in self.polymap.ridge_vertices:
self.corner_graph.add_edge(*r)
def feature(self, x, y):
""":param x, y 2-D coordinate."""
dist, idx = self.dist_map.query([x, y])
return self.poly_graph.node[idx]['feature']
def kdtree(ref_ra,ref_dec, ra,dec, k=1, dsmax=1./3600,verb=True):
'''computes distance (deg separation)= sqrt{(dec1-dec2)^2+ cos[0.5*(dec1+dec2)]^2*(ra1-ra2)^2}
return two dicts:
1)indices of reference data where matched and unmatched
2)indices of comparison data where matched, unmatched
UNITS: degrees for all
'''
print "use johan_tree() instead, for certain ra,dec finds <<< sources than astrometry spherematch"
raise ValueError
assert(len(ref_ra) == len(ref_dec))
assert(len(ra) == len(dec))
#index tree
tree = KDTree(np.transpose([dec.copy(),np.cos(dec.copy()*np.pi/180)*ra.copy()])) #data has shape NxK, N points and K dimensions (K=2 for ra,dec)
#look in tree for each references point, return k=1 NN
ds, i_tree = tree.query(np.transpose([ref_dec.copy(),np.cos(ref_dec.copy()*np.pi/180)*ref_ra.copy()]), k=k) #id for each data.query source that is NN of each input source
#make sense of results
i_ref,i_other={},{}
i_ref['match']= np.arange(len(ref_ra))[ds<=dsmax]
i_ref['nomatch']= np.arange(len(ref_ra))[ds>dsmax]
i_other['match']= i_tree[ds<=dsmax]
i_other['nomatch']= i_tree[ds>dsmax]
assert(len(ref_ra[i_ref['match']]) == len(ra[i_other['match']]))
if verb: print "%d/%d Refs matched, %d/%d Refs unmatched" % \
(i_ref['match'].size,len(ref_ra),i_ref['nomatch'].size,len(ref_ra))
return np.array(i_ref['match']),np.array(i_other['match']),np.array(ds[ds<=dsmax])
def _pointwiseF(self):
d = np.array(self.rcentroids)
tree = KDTree(d)
# get the 4 nearest neighbors (obviously the first one is the point itself)
nn = tree.query(d,k=4)[1]
N = len(self.rcentroids)
X = []
x = []
for i in xrange(N):
W = np.zeros((3,3),float)
w = np.zeros((3,3),float)
for j in xrange(3):
W[j,:] = self.rcentroids[nn[i,j+1]]-self.rcentroids[i]
w[j,:] = self.dcentroids[nn[i,j+1]]-self.dcentroids[i]
X.append(W)
x.append(w)
#Solve the linear system for each pointwise F
for i in xrange(N):
W = X[i]
w = x[i]
A = np.zeros((9,9),float)
A[0,0:3] = W[0,:]
A[1,3:6] = W[0,:]
A[2,6:9] = W[0,:]
A[3,0:3] = W[1,:]
A[4,3:6] = W[1,:]
A[5,6:9] = W[1,:]
A[6,0:3] = W[2,:]
A[7,3:6] = W[2,:]
A[8,6:9] = W[2,:]
b = w.flatten().T
F = np.linalg.solve(A,b)
self.localFs.append(F.reshape(3,3))
def main(in_file: Path, target_grid_file: Path, out_dir: Path=None):
if out_dir is not None:
out_dir.mkdir(exist_ok=True)
out_file = out_dir / (in_file.name + "_interpolated")
else:
out_file = in_file.parent / (in_file.name + "_interpolated")
if out_file.exists():
print(f"Skipping {in_file}, output already exists ({out_file})")
return
with xarray.open_dataset(target_grid_file) as ds_grid:
lons, lats = ds_grid["lon"][:].values, ds_grid["lat"][:].values
xt, yt, zt = lat_lon.lon_lat_to_cartesian(lons.flatten(), lats.flatten())
with xarray.open_dataset(in_file) as ds_in:
lons_s, lats_s = ds_in["lon"][:].values, ds_in["lat"][:].values
xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons_s.flatten(), lats_s.flatten())
ktree = KDTree(list(zip(xs, ys, zs)))
dists, inds = ktree.query(list(zip(xt, yt, zt)), k=1)
# resample to daily
ds_in_r = ds_in.resample(t="1D", keep_attrs=True).mean()
ds_out = xarray.Dataset()
for vname, var in ds_grid.variables.items():
ds_out[vname] = var[:]
ds_out["t"] = ds_in_r["t"][:]
for vname, var in ds_in_r.variables.items():
assert isinstance(var, xarray.Variable)
var = var.squeeze()
# only interested in (t, x, y) fields
if var.ndim != 3:
print(f"skipping {vname}")
continue
if vname.lower() not in ["t", "time", "lon", "lat"]:
print(f"Processing {vname}")
var_interpolated = [var[ti].values.flatten()[inds].reshape(lons.shape) for ti in range(var.shape[0])]
ds_out[vname] = xarray.DataArray(
var_interpolated, dims=("t", "x", "y"),
attrs=var.attrs,
)
ds_out.to_netcdf(out_file)
class AckermannController:
def __init__(self):
rospy.init_node("controller", anonymous=True)
self.time = rospy.get_time()
# retrieve ros parameters
self.max_speed = rospy.get_param("~max_speed")
self.time_step = rospy.get_param("~time_step")
self.traj_steps = rospy.get_param("~plan_steps")
ctrl_freq = rospy.get_param("~ctrl_freq")
rolename = rospy.get_param("~rolename")
self.steer_prev = 0.0
# state information
self.stateReady = False
self.state = odom_state()
# path tracking information
self.pathReady = False
self.path = np.zeros(shape=(self.traj_steps, 2)) # absolute position
self.path_tree = KDTree(self.path)
self.vel_path = np.zeros(shape=(self.traj_steps, 2))
self.steer_cache = None
# # PID controller parameter
self.pid_str_prop = rospy.get_param("~str_prop")
self.pid_str_deriv = rospy.get_param("~str_deriv")
# subscribers, publishers
rospy.Subscriber("/carla/" + rolename + "/odometry", Odometry,
self.odom_cb)
rospy.Subscriber("MSLcar0/command/trajectory", MultiDOFJointTrajectory,
self.desired_waypoints_cb)
self.command_pub = rospy.Publisher("/carla/" + rolename +
"/ackermann_cmd",
AckermannDrive,
queue_size=10)
# self.vehicle_cmd_pub = rospy.Publisher(
# "/carla/" + rolename + "/vehicle_control_cmd",
# CarlaEgoVehicleControl,
# queue_size=10
# )
self.tracking_pt_viz_pub = rospy.Publisher("tracking_point_mkr",
Marker,
queue_size=10)
self.ctrl_timer = rospy.Timer(rospy.Duration(1.0 / ctrl_freq),
self.timer_cb)
def desired_waypoints_cb(self, msg):
for i, pt in enumerate(msg.points):
self.path[i, 0] = pt.transforms[0].translation.x
self.path[i, 1] = pt.transforms[0].translation.y
self.vel_path[i, 0] = pt.velocities[0].linear.x
self.vel_path[i, 1] = pt.velocities[0].linear.y
self.path_tree = KDTree(self.path)
if not self.pathReady:
self.pathReady = True
def odom_cb(self, msg):
self.state.update_vehicle_state(msg)
if not self.stateReady:
self.stateReady = True
def timer_cb(self, event):
cmd_msg = AckermannDrive()
current_time = rospy.get_time()
# delta_t = current_time - self.time
delta_t = self.time_step
# delta_t = 0.05
self.time = current_time
jerk = 2.0
# cmd_msg.header.stamp = rospy.Time.now()
cmd_msg.steering_angle = 0.0
cmd_msg.steering_angle_velocity = 0.0
cmd_msg.speed = 0.0
cmd_msg.acceleration = 0.0
cmd_msg.jerk = jerk
if self.pathReady and self.stateReady:
pos_x, pos_y = self.state.get_position()
# pick target w/o collision avoidance, closest point on the traj and one point ahead
_, idx = self.path_tree.query([pos_x, pos_y])
# Target point for steering
if idx < self.traj_steps - 3:
target_pt = self.path_tree.data[idx + 3, :]
# str_idx = idx + 2
else:
target_pt = self.path_tree.data[-1, :]
# str_idx = self.vel_path.shape[0] - 1
print("CONTROLLER: at the end of the desired waypoits!!!")
# Target point for velocity
if idx < self.traj_steps:
target_vel = self.vel_path[idx, :]
# str_idx = idx + 2
else:
target_vel = self.vel_path[-1, :]
# str_idx = self.vel_path.shape[0] - 1
target_speed = np.linalg.norm(target_vel)
speed_diff = target_speed - self.state.get_speed()
# pos_diff = target_pt - self.state.get_position()
# acceleration = cmd_msg.acceleration + 2.0*(
# pos_diff/self.time_step**2.0 - speed_diff/self.time_step
# )
acceleration = abs(speed_diff) / (2.0 * delta_t)
cmd_msg.acceleration = np.min([1.5, acceleration])
steer = self.compute_ackermann_steer(target_pt)
steer_diff = abs(steer - self.steer_prev)
if steer_diff >= 0.3:
print(" 0.3")
steer = self.steer_prev
elif steer_diff >= 0.05:
print(" 0.05")
acceleration = 0.0
target_speed = target_speed / 5.0
self.steer_prev = steer
if self.state.get_speed() - target_speed > 0.0:
# cmd_msg.acceleration = acceleration - 5
acceleration = -100.0
cmd_msg.speed = target_speed / 5.0
jerk = 1000.0
elif target_speed - self.state.get_speed() > 0.2:
acceleration = np.min([3.0, acceleration])
else:
acceleration = 0.0
cmd_msg.steering_angle = steer
cmd_msg.speed = target_speed
cmd_msg.acceleration
cmd_msg.jerk = jerk
# Print out some debugging information
if abs(target_speed - self.max_speed) > 2.0:
# print("Posit diff:", pos_diff)
print("Speed diff:", speed_diff)
print("Current speed:", self.state.get_speed())
print("Desired speed:", target_speed)
print("Desired accel:", acceleration)
print("Desired jerk:", jerk)
print("Speed (sent):", cmd_msg.speed)
print("Accel (sent):", cmd_msg.acceleration)
print("Jerk (sent):", cmd_msg.jerk)
print("delta t:", delta_t)
# for visualization purposes and debuging control node
mk_msg = Marker()
mk_msg.header.stamp = rospy.Time.now()
mk_msg.header.frame_id = 'map'
mk_msg.pose.position.x = target_pt[0]
mk_msg.pose.position.y = target_pt[1]
mk_msg.type = Marker.CUBE
mk_msg.scale.x = 3
mk_msg.scale.y = 3
mk_msg.scale.z = 3
mk_msg.color.a = 1.0
mk_msg.color.r = 1
mk_msg.color.b = 1
self.tracking_pt_viz_pub.publish(mk_msg)
# self.vehicle_cmd_pub.publish(vehicle_cmd_msg)
self.command_pub.publish(cmd_msg)
def compute_ackermann_steer(self, target_pt):
pos_x, pos_y, yaw = self.state.get_pose()
if np.linalg.norm([target_pt[0] - pos_x, target_pt[1] - pos_y]) < 1:
print("target point too close!!!!!!!!")
if self.steer_cache:
return self.steer_cache
else:
return 0.0
egoOri = np.array([np.cos(yaw), np.sin(yaw), 0])
rel_pos = np.array([target_pt[0] - pos_x, target_pt[1] - pos_y, 0])
rel_pos_norm = np.linalg.norm(rel_pos)
rel_pos_unit = rel_pos / np.linalg.norm(rel_pos)
rot = np.cross(rel_pos_unit, egoOri)
steer = -rot[2] * self.pid_str_prop
if self.steer_cache is not None:
d_steer = (steer - self.steer_cache)
else:
d_steer = 0.0
# print("steer, d_steer:", steer, d_steer)
steer = steer + d_steer * self.pid_str_deriv
self.steer_cache = steer
return steer
class Partitioner(object):
"""
Universe of Discourse partitioner. Split data on several fuzzy sets
"""
def __init__(self, **kwargs):
"""
Universe of Discourse partitioner scheme. Split data on several fuzzy sets
"""
self.name = kwargs.get('name', "")
"""partitioner name"""
self.partitions = kwargs.get('npart', 10)
"""The number of universe of discourse partitions, i.e., the number of fuzzy sets that will be created"""
self.sets = {}
self.membership_function = kwargs.get('func', Membership.trimf)
"""Fuzzy membership function (pyFTS.common.Membership)"""
self.setnames = kwargs.get('names', None)
"""list of partitions names. If None is given the partitions will be auto named with prefix"""
self.prefix = kwargs.get('prefix', 'A')
"""prefix of auto generated partition names"""
self.transformation = kwargs.get('transformation', None)
"""data transformation to be applied on data"""
self.indexer = kwargs.get('indexer', None)
self.variable = kwargs.get('variable', None)
"""In a multivariate context, the variable that contains this partitioner"""
self.type = kwargs.get('type', 'common')
"""The type of fuzzy sets that are generated by this partitioner"""
self.ordered_sets = None
"""A ordered list of the fuzzy sets names, sorted by their middle point"""
self.kdtree = None
"""A spatial index to help in fuzzyfication"""
if kwargs.get('preprocess', True):
data = kwargs.get('data', [None])
if self.indexer is not None:
ndata = self.indexer.get_data(data)
else:
ndata = data
if self.transformation is not None:
ndata = self.transformation.apply(ndata)
else:
ndata = data
if self.indexer is not None:
ndata = self.indexer.get_data(ndata)
_min = np.nanmin(ndata)
if _min == -np.inf:
ndata[ndata == -np.inf] = 0
_min = np.nanmin(ndata)
self.min = float(_min * 1.1 if _min < 0 else _min * 0.9)
_max = np.nanmax(ndata)
self.max = float(_max * 1.1 if _max > 0 else _max * 0.9)
self.sets = self.build(ndata)
self.partitions = len(self.sets)
if self.ordered_sets is None and self.setnames is not None:
self.ordered_sets = self.setnames[:len(self.sets)]
else:
self.ordered_sets = FuzzySet.set_ordered(self.sets)
del (ndata)
def extractor(self, x):
"""Extract a single primitive type from an structured instance"""
return x
def build(self, data):
"""
Perform the partitioning of the Universe of Discourse
:param data: training data
:return:
"""
pass
def get_name(self, counter):
"""
Find the name of the fuzzy set given its counter id.
:param counter: The number of the fuzzy set
:return: String
"""
return self.prefix + str(
counter) if self.setnames is None else self.setnames[counter]
def lower_set(self):
"""
Return the fuzzy set on lower bound of the universe of discourse.
:return: Fuzzy Set
"""
return self.sets[self.ordered_sets[0]]
def upper_set(self):
"""
Return the fuzzy set on upper bound of the universe of discourse.
:return: Fuzzy Set
"""
return self.sets[self.ordered_sets[-1]]
def build_index(self):
points = []
#self.index = {}
for ct, key in enumerate(self.ordered_sets):
fset = self.sets[key]
points.append([fset.lower, fset.centroid, fset.upper])
#self.index[ct] = fset.name
import sys
sys.setrecursionlimit(100000)
self.kdtree = KDTree(points)
sys.setrecursionlimit(1000)
def fuzzyfy(self, data, **kwargs):
"""
Fuzzyfy the input data according to this partitioner fuzzy sets.
:param data: input value to be fuzzyfied
:keyword alpha_cut: the minimal membership value to be considered on fuzzyfication (only for mode='sets')
:keyword method: the fuzzyfication method (fuzzy: all fuzzy memberships, maximum: only the maximum membership)
:keyword mode: the fuzzyfication mode (sets: return the fuzzy sets names, vector: return a vector with the membership
values for all fuzzy sets, both: return a list with tuples (fuzzy set, membership value) )
:returns a list with the fuzzyfied values, depending on the mode
"""
if isinstance(data, (tuple, list, np.ndarray)):
ret = []
for inst in data:
mv = self.fuzzyfy(inst, **kwargs)
ret.append(mv)
return ret
alpha_cut = kwargs.get('alpha_cut', 0.)
mode = kwargs.get('mode', 'sets')
method = kwargs.get('method', 'fuzzy')
nearest = self.search(data, type='index')
mv = np.zeros(self.partitions)
try:
for ix in nearest:
tmp = self[ix].membership(data)
mv[ix] = tmp if tmp >= alpha_cut else 0.
except:
print(ix)
ix = np.ravel(np.argwhere(mv > 0.))
if ix.size == 0:
mv[self.check_bounds(data)] = 1.
if method == 'fuzzy' and mode == 'vector':
return mv
elif method == 'fuzzy' and mode == 'sets':
try:
ix = np.ravel(np.argwhere(mv > 0.))
sets = [
self.ordered_sets[i] for i in ix if i < self.partitions
]
return sets
except Exception as ex:
return None
elif method == 'maximum' and mode == 'sets':
mx = max(mv)
ix = np.ravel(np.argwhere(mv == mx))
return self.ordered_sets[ix[0]]
elif mode == 'both':
ix = np.ravel(np.argwhere(mv > 0.))
sets = [(self.ordered_sets[i], mv[i]) for i in ix]
return sets
def defuzzyfy(self, values, mode='both'):
if not isinstance(values, list):
values = [values]
num = []
den = []
for val in values:
fset = val[0]
mv = val[1]
if mode == 'both':
num.append(self.sets[fset].centroid * mv)
den.append(mv)
elif mode == 'sets':
num.append(self.sets[fset].centroid)
elif mode == 'vector':
num.append(self.sets[self.ordered_sets[fset]].centroid * mv)
den.append(mv)
else:
raise Exception('Unknown deffuzyfication mode')
if mode in ('both', 'vector'):
return np.nansum(num) / np.nansum(den)
else:
return np.nanmean(num)
def check_bounds(self, data):
"""
Check if the input data is outside the known Universe of Discourse and, if it is, round it to the closest
fuzzy set.
:param data: input data to be verified
:return: the index of the closest fuzzy set when data is outside de universe of discourse or None if
the data is inside the UoD.
"""
if data < self.min:
return 0
elif data > self.max:
return self.partitions - 1
def search(self, data, **kwargs):
"""
Perform a search for the nearest fuzzy sets of the point 'data'. This function were designed to work with several
overlapped fuzzy sets.
:param data: the value to search for the nearest fuzzy sets
:param type: the return type: 'index' for the fuzzy set indexes or 'name' for fuzzy set names.
:param results: the number of nearest fuzzy sets to return
:return: a list with the nearest fuzzy sets
"""
if self.kdtree is None:
self.build_index()
type = kwargs.get('type', 'index')
results = kwargs.get('results', 3)
_, ix = self.kdtree.query([data, data, data], results)
if type == 'name':
return [self.ordered_sets[k] for k in sorted(ix)]
else:
return sorted(ix)
def plot(self, ax, rounding=0):
"""
Plot the partitioning using the Matplotlib axis ax
:param ax: Matplotlib axis
"""
ax.set_title(self.name)
ax.set_ylim([0, 1.1])
ax.set_xlim([self.min, self.max])
ticks = []
x = []
for key in self.sets.keys():
s = self.sets[key]
if s.type == 'common':
self.plot_set(ax, s)
elif s.type == 'composite':
for ss in s.sets:
self.plot_set(ax, ss)
ticks.append(str(round(s.centroid, rounding)) + '\n' + s.name)
x.append(s.centroid)
ax.xaxis.set_ticklabels(ticks)
ax.xaxis.set_ticks(x)
def plot_set(self, ax, s):
"""
Plot an isolate fuzzy set on Matplotlib axis
:param ax: Matplotlib axis
:param s: Fuzzy Set
"""
if s.mf == Membership.trimf:
ax.plot([s.parameters[0], s.parameters[1], s.parameters[2]],
[0, s.alpha, 0])
elif s.mf in (Membership.gaussmf, Membership.bellmf, Membership.sigmf):
tmpx = np.linspace(s.lower, s.upper, 100)
tmpy = [s.membership(kk) for kk in tmpx]
ax.plot(tmpx, tmpy)
elif s.mf == Membership.trapmf:
ax.plot(s.parameters, [0, s.alpha, s.alpha, 0])
elif s.mf == Membership.singleton:
ax.plot([s.parameters[0], s.parameters[0]], [0, s.alpha])
def __str__(self):
"""
Return a string representation of the partitioner, the list of fuzzy sets and their parameters
:return:
"""
tmp = self.name + ":\n"
for key in self.sets.keys():
tmp += str(self.sets[key]) + "\n"
return tmp
def __len__(self):
"""
Return the number of partitions
:return: number of partitions
"""
return self.partitions
def __getitem__(self, item):
"""
Return a fuzzy set by its order or its name.
:param item: If item is an integer then it represents the fuzzy set index (order), if it was a string then
it represents the fuzzy set name.
:return: the fuzzy set
"""
if isinstance(item,
(int, np.int, np.int8, np.int16, np.int32, np.int64)):
if item < 0 or item >= self.partitions:
raise ValueError(
"The fuzzy set index must be between 0 and {}.".format(
self.partitions))
return self.sets[self.ordered_sets[item]]
elif isinstance(item, str):
if item not in self.sets:
raise ValueError(
"The fuzzy set with name {} does not exist.".format(item))
return self.sets[item]
else:
raise ValueError(
"The parameter 'item' must be an integer or a string and the value informed was {} of type {}!"
.format(item, type(item)))
def __iter__(self):
"""
Iterate over the fuzzy sets, ordered by its midpoints.
:return: An iterator over the fuzzy sets.
"""
for key in self.ordered_sets:
yield self.sets[key]
def estimate_overall_ref_model(colmap_folder, min_track_len, min_tri_angle,
max_tri_angle, image_path):
print 'Loading COLMAP data'
scene_manager = SceneManager(colmap_folder)
scene_manager.load_cameras()
scene_manager.load_images()
images = [
'frame0859.jpg', 'frame0867.jpg', 'frame0875.jpg', 'frame0883.jpg',
'frame0891.jpg'
]
_L = np.array([])
_r = np.array([])
_ndotl = np.array([])
for image_name in images:
image_id = scene_manager.get_image_id_from_name(image_name)
image = scene_manager.images[image_id]
camera = scene_manager.get_camera(image.camera_id)
# image pose
R = util.quaternion_to_rotation_matrix(image.qvec)
print 'Loading 3D points'
scene_manager.load_points3D()
scene_manager.filter_points3D(min_track_len,
min_tri_angle=min_tri_angle,
max_tri_angle=max_tri_angle,
image_list=set([image_id]))
points3D, points2D = scene_manager.get_points3D(image_id)
points3D = points3D.dot(R.T) + image.tvec[np.newaxis, :]
# need to remove redundant points
# http://stackoverflow.com/questions/16970982/find-unique-rows-in-numpy-array
points2D_view = np.ascontiguousarray(points2D).view(
np.dtype((np.void, points2D.dtype.itemsize * points2D.shape[1])))
_, idx = np.unique(points2D_view, return_index=True)
points2D, points3D = points2D[idx], points3D[idx]
# further rule out any points too close to the image border (for bilinear
# interpolation)
mask = ((points2D[:, 0] < camera.width - 1) &
(points2D[:, 1] < camera.height - 1))
points2D, points3D = points2D[mask], points3D[mask]
points2D_image = points2D.copy() # coordinates in image space
points2D = np.hstack((points2D, np.ones((points2D.shape[0], 1))))
points2D = points2D.dot(np.linalg.inv(
camera.get_camera_matrix()).T)[:, :2]
print len(points3D), 'total points'
# load image
#image_file = scene_manager.image_path + image.name
image_file = image_path + image.name
im_rgb = skimage.img_as_float(
skimage.io.imread(image_file)) # color image
L = skimage.color.rgb2lab(im_rgb)[:, :, 0] * 0.01
L = np.maximum(
L, 1e-6) # unfortunately, can't have black pixels, since we
# divide by L
# initial values on unit sphere
x, y = camera.get_image_grid()
print 'Computing nearest neighbors'
kdtree = KDTree(points2D)
weights, nn_idxs = kdtree.query(np.c_[x.ravel(), y.ravel()],
KD_TREE_KNN)
weights = weights.reshape(camera.height, camera.width, KD_TREE_KNN)
nn_idxs = nn_idxs.reshape(camera.height, camera.width, KD_TREE_KNN)
# turn distances into weights for the nearest neighbors
np.exp(-weights, weights) # in-place
# create initial surface on unit sphere
S0 = np.dstack((x, y, np.ones_like(x)))
S0 /= np.linalg.norm(S0, axis=-1)[:, :, np.newaxis]
r_fixed = np.linalg.norm(points3D, axis=-1) # fixed 3D depths
S = S0.copy()
z = S0[:, :, 2]
S, r_ratios = nearest_neighbor_warp(weights, nn_idxs, points2D_image,
r_fixed,
util.generate_surface(camera, z))
z_est = np.maximum(S[:, :, 2], 1e-6)
S = util.generate_surface(camera, z_est)
r = np.linalg.norm(S, axis=-1)
ndotl = util.calculate_ndotl(camera, S)
if _L.size == 0:
_L = L.ravel()
_r = r.ravel()
_ndotl = ndotl.ravel()
else:
_L = np.append(_L, L.ravel())
_r = np.append(_r, r.ravel())
_ndotl = np.append(_ndotl, ndotl.ravel())
#
falloff, model_params, residual = fit_reflectance_model(
sfs.POWER_MODEL, _L, _r, _ndotl, False, 5)
import pdb
pdb.set_trace()
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.camera_image = None
self.lights = []
self.waypoints_2d = None
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
sub4 = rospy.Subscriber('/image_color', Image, self.image_cb)
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
self.sim = self.config["sim"]
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.bridge = CvBridge()
self.listener = tf.TransformListener()
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
self.waypoint_tree = None
self.debug = self.config["debug"]
if self.debug:
self.debug_file = open(os.path.join(os.getcwd(), 'debug.csv'), 'w')
self.debug_file.write(
'Image,Prediction,Ground Truth,Prediction Time (sec)\n')
self.light_classifier = TLClassifier(self.sim)
rospy.spin()
#self.loop()
def loop(self):
rate = rospy.Rate(ROS_RATE)
while not rospy.is_shutdown():
#self.image_cb_dummy()############################################################
rate.sleep()
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, waypoints):
self.waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in waypoints.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d)
def traffic_cb(self, msg):
self.lights = msg.lights
#rospy.logwarn('traffic_cb called!!!')
def image_cb_dummy(
self
): #####################################################################
self.image_cb(None)
def image_cb(self, msg):
"""Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.has_image = True
self.camera_image = msg
#rospy.logwarn('Image dims --- Height: %s, Width: %s', msg.height, msg.width)
#image = self.bridge.imgmsg_to_cv2(msg, 'rgb8')
#cv2.imwrite('/home/student/CarND-Capstone-master/imgs/simulator/' + str(rospy.Time.now()) + '.png', image)
light_wp, state = self.process_traffic_lights()
'''
Publish upcoming red lights at camera frequency.
Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
of times till we start using it. Otherwise the previous stable state is
used.
'''
if self.state != state:
self.state_count = 0
self.state = state
elif self.state_count >= STATE_COUNT_THRESHOLD:
self.last_state = self.state
light_wp = light_wp if state == TrafficLight.RED else -1
self.last_wp = light_wp
self.upcoming_red_light_pub.publish(Int32(light_wp))
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
self.state_count += 1
def get_closest_waypoint(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
#TODO implement
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
return closest_idx
def get_light_state(self, light):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
if (not self.has_image):
self.prev_light_loc = None
return False
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "rgb8")
#Get classification
if self.debug:
start_time = time.time()
result = self.light_classifier.get_classification(cv_image)
if self.debug:
end_time = time.time()
elapsed = end_time - start_time
image_name = str(rospy.Time.now()) + '-' + str(result) + '.png'
#img_path = os.getcwd() + '/../../../imgs/' + image_name
#cv2.imwrite(img_path, cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB))
self.debug_file.write('{},{},{},{:.2f}\n'.format(
image_name, result, light.state, elapsed))
return result
# ############################################### TESTING ONLY ##################################################################
#return light.state
def process_traffic_lights(self):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
closest_light = None
line_wp_idx = None
# List of positions that correspond to the line to stop in front of for a given intersection
stop_line_positions = self.config['stop_line_positions']
if (self.pose):
car_wp_idx = self.get_closest_waypoint(self.pose.pose.position.x,
self.pose.pose.position.y)
#TODO find the closest visible traffic light (if one exists)
diff = len(self.waypoints.waypoints)
for i, light in enumerate(self.lights):
# Get stop line waypoint index
line = stop_line_positions[i]
temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
# Get closest stop line waypoint index
d = temp_wp_idx - car_wp_idx
if d >= 0 and d < diff:
diff = d
closest_light = light
line_wp_idx = temp_wp_idx
if closest_light:
state = self.get_light_state(closest_light)
#rospy.logwarn('Traffic light state: %s', state)
#rospy.logwarn('Closest traffic light stop line index: %s, state: %s', line_wp_idx, state)
return line_wp_idx, state
#self.waypoints = None
return -1, TrafficLight.UNKNOWN
def __del__(self):
if self.debug:
self.debug_file.close()
class face_manager:
def __init__(self):
self.name_to_face_vec_list = {}
self.name_to_face_img_list = {}
'''
for example:
username1:[face_vec1,face_vec2,face_vec3,....]
'''
#self.pca_processed_name_to_face_vec_list = {}
self.kd_tree = None
self.vec_to_name_dict = {}
self.vec_list = None
self.name_list = None
self.similar_threshold = SIMILAR_THRES
def build_kdtree(self):
#step1 build inversed mapping.
self.vec_list = []
self.name_list = []
for name in self.name_to_face_vec_list:
for vec in self.name_to_face_vec_list[name]:
self.vec_to_name_dict[vec] = name
self.vec_list.append(np.array(vec))
self.name_list.append(name)
vec_mat = np.array(self.vec_list)
self.kd_tree = KDTree(vec_mat)
def add_name_and_img(self,name,cv_img_bgr_origin):
from Network.facenet.get_face_vec import get_vec_from_cvimg
# if self.name_to_face_img_list.has_key(name):
if name in self.name_to_face_img_list.keys():
self.name_to_face_img_list[name].append(cv_img_bgr_origin)
self.name_to_face_vec_list[name].append(get_vec_from_cvimg(cv_img_bgr_origin))
else:
self.name_to_face_img_list[name] = [cv_img_bgr_origin]
self.name_to_face_vec_list[name] = [get_vec_from_cvimg(cv_img_bgr_origin)]
def query_name_by_img(self,cv_img_bgr_origin):
if self.kd_tree is None:
self.build_kdtree()
from Network.facenet.get_face_vec import get_vec_from_cvimg
vec = get_vec_from_cvimg(cv_img_bgr_origin)
similar_vecs = self.kd_tree.query_ball_point(vec,self.similar_threshold)
similar_vecs = map(lambda x:np.array(x),similar_vecs)
if len(similar_vecs)>0:
if len(similar_vecs) == 1:
return self.name_list[int(similar_vecs[0])] # one match.return it.
elif len(similar_vecs)>1:
min_dist,min_id = self.similar_threshold+1,0 # ensures that any of them smaller than THRES+1.
for i in range(len(similar_vecs)):
dist = np.sqrt(np.sum(np.square(np.array(similar_vecs[i]) - np.array(vec))))
if dist <min_dist:
min_dist = dist
min_id = i
print ('similar_vecs',similar_vecs)
return self.name_list[int(similar_vecs[min_id])]
else:
return None # no matching.
def query_nearest_by_img(self,cv_img_bgr_origin,return_score=False):
if self.kd_tree is None:
self.build_kdtree()
from Network.facenet.get_face_vec import get_vec_from_cvimg
time0 = time.time()
vec = get_vec_from_cvimg(cv_img_bgr_origin)
time1 = time.time()
dist,nearest_vec_id = self.kd_tree.query([vec])
time2 = time.time()
print ('vec calc time cost:',time1-time0)
print ('kdtree query time cost:',time2-time1)
dist,nearest_vec_id = dist[0],nearest_vec_id[0]
print ('dist square:',dist*dist)
if dist*dist < self.similar_threshold:
if return_score:
return self.name_list[nearest_vec_id],dist*dist
else:
return self.name_list[nearest_vec_id]
else:
if return_score:
return None,None
else:
return None
def method12(dists, N):
tree = KDTree(dists, leafsize=dists.shape[0] + 1)
distances, ndx = tree.query(0, k=N)
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.waypoints_2d = None
self.waypoint_tree = None
self.stateCount = 0
self.image_count = 0
self.state = TrafficLight.RED
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
self.camera_image = None
self.lights = []
self.saveImgs = True
if self.saveImgs:
if not (os.path.exists("./saveImgs")):
os.mkdir("./saveImgs")
self.saveCount = 0
config_string = rospy.get_param("/traffic_light_config")
#config_string2 = config_string
#rospy.loginfo("+++++++++++++++Using simulator+++++++++++++++%s",config_string)
self.config = yaml.load(config_string)
#rospy.loginfo("+++++++++++++++Using simulator+++++++++++++++%s",self.config)
isSite = bool(rospy.get_param("~is_siteP", True))
if isSite:
self.usingSimulator = False
rospy.loginfo("+++++++++++++++Using simulator+++++++++++++++")
else:
self.usingSimulator = True
rospy.loginfo("+++++++++++++++Using simulator+++++++++++++++")
self.usingSystemLightState = 0
#self.usingSimulator = 0 if self.config['is_site'] else 1
#self.usingSimulator = bool(rospy.get_param("~is_siteP", False))
self.bridge = CvBridge()
self.light_classifier = TLClassifier(self.usingSimulator)
self.listener = tf.TransformListener()
self.stop_closest_waypoint = []
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
#sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)
sub6 = rospy.Subscriber('/image_color',
Image,
self.image_cb,
queue_size=1,
buff_size=2 * 52428800)
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.lastTime = rospy.get_time()
rate = rospy.Rate(200)
while not rospy.is_shutdown():
if self.pose and self.waypoints and self.lights:
#get closest waypoint
#closest_waypoint_idx = self.get_closest_waypoint_idx()
#rospy.loginfo('closest_waypoint_index:%s', closest_waypoint_idx)
#self.publish_waypoints(closest_waypoint_idx)
self.InitializeImage = True
light_wp, state = self.process_traffic_lights()
self.find_traffic_lights(light_wp, state)
rospy.loginfo(
"=============finish initialize image===========")
self.InitializeImage = False
break
rate.sleep()
rospy.spin()
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, waypoints):
self.waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in waypoints.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d)
def traffic_cb(self, msg):
self.lights = msg.lights
def image_cb(self, msg):
"""Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.has_image = True
ThisTime = rospy.get_time()
#rospy.loginfo("Time elapsed:%s",ThisTime - self.lastTime)
self.lastTime = ThisTime
self.image_count = self.image_count + 1
THRESHOLD_SAMPLE = 1
#if self.usingSimulator:
#THRESHOLD_SAMPLE = 1
if (self.image_count >= THRESHOLD_SAMPLE):
self.lastTime = rospy.get_time()
self.image_count = 0
self.has_image = True
self.camera_image = msg
light_wp, state = self.process_traffic_lights()
if (state != 0) and self.saveImgs:
iimage = self.bridge.imgmsg_to_cv2(self.camera_image, "rgb8")
h, w, _ = iimage.shape
#rospy.loginfo("image width:%s height:%s state:%s",w,h,state)
if self.usingSimulator:
iimage = iimage[0:int(0.7 * h), 0:w]
else:
iimage = iimage[0:int(0.7 * h), 50:int(w - 50)]
self.saveImags(iimage, state)
ThisTime = rospy.get_time()
#rospy.loginfo("Time elapsed:%s, light state:%s",ThisTime - self.lastTime,state)
self.find_traffic_lights(light_wp, state)
def saveImags(self, image, state):
dictTL = {0: "R", 1: "Y", 2: "G", 4: "U"}
takeImage = image
if not self.usingSimulator:
lsImageName = "./saveImgs/image0{0:0>5}.jpg".format(self.saveCount)
#rospy.loginfo("save image:%s",lsImageName)
cv2.imwrite(lsImageName, takeImage)
else:
lsImageName = "./saveImgs/{0}_image6{1:0>5}.jpg".format(
dictTL[state], self.saveCount)
rospy.loginfo("save image:%s", lsImageName)
cv2.imwrite(lsImageName, takeImage)
self.saveCount += 1
def find_traffic_lights(self, light_wp, state):
'''
Publish upcoming red lights at camera frequency.
Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
of times till we start using it. Otherwise the previous stable state is
used.
'''
if state == TrafficLight.YELLOW:
state = TrafficLight.RED
if self.InitializeImage:
self.last_wp = light_wp
self.upcoming_red_light_pub.publish(Int32(light_wp))
return
if self.state != state:
self.state_count = 0
self.state = state
elif self.state_count >= STATE_COUNT_THRESHOLD:
self.last_state = self.state
light_wp = light_wp if state == TrafficLight.RED else -1
self.last_wp = light_wp
rospy.loginfo("---light changed,wp is:%s state:%s s state:%s",
light_wp, state, self.state)
self.upcoming_red_light_pub.publish(Int32(light_wp))
else:
rospy.loginfo("---light remained,wp is:%s state:%s s state:%s",
self.last_wp, state, self.state)
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
self.state_count += 1
def get_closest_waypoint(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
#TODO implement
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
# check if the closest is ahed or behind the vehicle
closest_coord = self.waypoints_2d[closest_idx]
prev_coord = self.waypoints_2d[closest_idx - 1]
closest_v = np.array(closest_coord)
prev_v = np.array(prev_coord)
pos_v = np.array([x, y])
val = np.dot(closest_v - prev_v, pos_v - closest_v)
if val > 0:
closest_idx = (closest_idx + 1) % len(self.waypoints_2d)
return closest_idx
def get_light_state(self, light):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
if self.usingSystemLightState > 0:
if (self.stateCount > 2):
#rospy.loginfo("light state:{0}".format(light.state))
self.stateCount = 0
self.stateCount = self.stateCount + 1
return light.state
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "rgb8")
#cv_image = cv2.cvtColor(cv_image, cv2.COLOR_BGR2RGB)
#narrow down seaching area of the image
h, w, _ = cv_image.shape
if self.usingSimulator:
cv_image = cv_image[0:int(0.7 * h), 0:w]
else:
cv_image = cv_image[0:int(0.7 * h), 0:w]
#Get classification
return self.light_classifier.get_classification(cv_image)
def process_traffic_lights(self):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
light = None
light_wp = None
# List of positions that correspond to the line to stop in front of for a given intersection
stop_line_positions = self.config['stop_line_positions']
if (self.pose):
car_position = self.get_closest_waypoint(self.pose.pose.position.x,
self.pose.pose.position.y)
#TODO find the closest visible traffic light (if one exists)
diff = 200
#for fast get
if len(self.stop_closest_waypoint) == 0:
for i, lightP in enumerate(self.lights):
line = stop_line_positions[i]
self.stop_closest_waypoint.append(
self.get_closest_waypoint(line[0], line[1]))
#rospy.loginfo("len of waypoints:%s car wp:%s",diff,car_position)
for i, lightP in enumerate(self.lights):
tmp_waypoint_idx = self.stop_closest_waypoint[i]
d = tmp_waypoint_idx - car_position
if d >= 0 and d < diff:
diff = d
light = lightP
light_wp = tmp_waypoint_idx
if light:
rospy.loginfo("car pos:%s closest light idx %s diff:%s",
car_position, light_wp, diff)
if self.InitializeImage:
#for first image latency
state = 0
else:
state = self.get_light_state(light)
return light_wp, state
#self.waypoints = None
return -1, TrafficLight.UNKNOWN
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.waypoints_2d = None
self.waypoints_tree = None
self.camera_image = None
self.lights = []
'''
The permanent "(x, y)" coordinates for each traffic light's stop line are provided by the
"config" dictionary, which is imported from the "/traffic_light_config" file
'''
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
'''
The node should publish the index of the waypoint for
nearest upcoming red light's stop line to a single topic is `/traffic_waypoint`
'''
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.bridge = CvBridge()
self.light_classifier = TLClassifier(self.config['is_site'])
self.listener = tf.TransformListener()
self.state = 4 # TrafficLight.UNKNOWN
# rospy.logwarn(">> __init__: self.state= {0}".format(self.state))
self.last_state = 4
self.last_wp = -1
self.state_count = 0
# `/current_pose` can be used to determine the vehicle's location
sub1 = rospy.Subscriber('/current_pose', PoseStamped,
self.current_pose_cb)
# `/base_waypoints` provides the complete list of waypoints for the course
sub2 = rospy.Subscriber('/base_waypoints', Lane,
self.base_waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.vehicle_traffic_lights_cb)
'''
`/image_color` which provides an image stream from the car's camera.
These images are used to determine the color of upcoming traffic lights
'''
# sub6 = rospy.Subscriber('/image_color', Image, self.image_color_cb, queue_size=1, buff_size=5760000)
sub6 = rospy.Subscriber('/image_color', Image, self.image_color_cb)
rospy.spin()
def current_pose_cb(self, msg):
self.pose = msg
def base_waypoints_cb(self, waypoints):
self.waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in waypoints.waypoints]
self.waypoints_tree = KDTree(self.waypoints_2d)
def vehicle_traffic_lights_cb(self, msg):
self.lights = msg.lights
def image_color_cb(self, msg):
"""
Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.has_image = True
self.camera_image = msg
light_wp, state = self.process_traffic_lights()
# rospy.logwarn("state: {0}".format(state))
# rospy.logwarn("self.state: {0}".format(self.state))
'''
Publish upcoming red lights at camera frequency.
Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
of times till we start using it. Otherwise the previous stable state is
used.
'''
if self.state != state:
self.state_count = 0
self.state = state
elif self.state_count >= STATE_COUNT_THRESHOLD:
self.last_state = self.state
if state == TrafficLight.RED:
light_wp = light_wp
elif state == TrafficLight.YELLOW:
light_wp = -light_wp
elif state == TrafficLight.GREEN:
light_wp = -2
else:
light_wp = -4
self.last_wp = light_wp
self.upcoming_red_light_pub.publish(Int32(light_wp))
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
self.state_count += 1
def get_closest_waypoint(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
x: x position to match a waypoint's "x" to
y: y position to match a waypoint's "y" to
Returns:
int: index of the closest waypoint in self.waypoints
"""
closest_idx = self.waypoints_tree.query([x, y], 1)[1]
return closest_idx
def get_light_state(self, light):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
# For testing just return light state
# return light.state
if (not self.has_image):
self.prev_light_loc = None
return False
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
#Get classification
return self.light_classifier.get_classification(cv_image)
def process_traffic_lights(self):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
closest_light = None
line_wp_idx = None
# List of positions that correspond to the line to stop in front of for a given intersection
stop_line_positions = self.config['stop_line_positions']
if (self.pose):
car_position_idx = self.get_closest_waypoint(
self.pose.pose.position.x, self.pose.pose.position.y)
#TODO find the closest visible traffic light (if one exists)
diff = len(self.waypoints.waypoints)
for i, light in enumerate(self.lights):
# Get stop line waypoint index
line = stop_line_positions[i]
temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
# Find closest stop line waypoint index
d = temp_wp_idx - car_position_idx
if d >= 0 and d < diff:
diff = d
closest_light = light
line_wp_idx = temp_wp_idx
# rospy.logwarn("process_traffic_lights | self.pose: {0}".format(self.pose))
if closest_light:
state = self.get_light_state(closest_light)
# rospy.logwarn("process_traffic_lights | {0}, {1}".format(line_wp_idx, state))
return line_wp_idx, state
# self.waypoints = None
# rospy.logwarn("process_traffic_lights | {0}, {1}".format(-1, TrafficLight.UNKNOWN))
return -1, TrafficLight.UNKNOWN
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.waypoints_2d = None
self.waypoint_tree = None
self.camera_image = None
self.lights = []
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.bridge = CvBridge()
self.light_classifier = TLClassifier()
self.listener = tf.TransformListener()
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
rospy.spin()
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, msg):
self.waypoints = msg.waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in self.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d)
def traffic_cb(self, msg):
self.lights = msg.lights
def image_cb(self, msg):
"""Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.has_image = True
self.camera_image = msg
light_wp, state = self.process_traffic_lights()
'''
Publish upcoming red lights at camera frequency.
Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
of times till we start using it. Otherwise the previous stable state is
used.
'''
if self.state != state:
self.state_count = 0
self.state = state
elif self.state_count >= STATE_COUNT_THRESHOLD:
self.last_state = self.state
light_wp = light_wp if state == TrafficLight.RED else -1
self.last_wp = light_wp
self.upcoming_red_light_pub.publish(Int32(light_wp))
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
self.state_count += 1
def get_closest_waypoint(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
current_postion_coord = [x, y]
closest_wp_idx = self.waypoint_tree.query(current_postion_coord, 1)[1]
return closest_wp_idx
def get_light_state(self, light):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
#return light.state
if (not self.has_image):
self.prev_light_loc = None
return False
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
#Get classification
return self.light_classifier.get_classification(cv_image)
def process_traffic_lights(self):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
light = None
closest_light_wp_idx = None
line_wp_index = None
# List of positions that correspond to the line to stop in front of for a given intersection
stop_line_positions = self.config['stop_line_positions']
if (self.pose):
car_position_idx = self.get_closest_waypoint(
self.pose.pose.position.x, self.pose.pose.position.y)
#TODO find the closest visible traffic light (if one exists)
diff = len(self.waypoints)
for i, light in enumerate(self.lights):
line = stop_line_positions[i]
temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
d = temp_wp_idx - car_position_idx
if d >= 0 and d < diff:
diff = d
closest_light_wp_idx = light
line_wp_index = temp_wp_idx
if closest_light_wp_idx:
state = self.get_light_state(closest_light_wp_idx)
return line_wp_index, state
return -1, TrafficLight.UNKNOWN
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.camera_image = None
self.lights = []
self.waypoints_2d = None
self.waypoints_tree = None
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
self.detector_config = {}
self.detector_config["model_anchor"] = rospy.get_param("model_anchor")
self.detector_config["model_classes"] = rospy.get_param(
"model_classes")
self.detector_config["model_path"] = rospy.get_param("model_path")
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.bridge = CvBridge()
self.light_classifier = TLClassifier(self.config['is_site'],
self.detector_config)
self.listener = tf.TransformListener()
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
self.prev_light_loc = None
self.has_image = False
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
if self.config['is_site']: # subscribe to raw image
sub6 = rospy.Subscriber('/image_raw',
Image,
self.image_cb,
queue_size=1,
tcp_nodelay=True)
else:
sub6 = rospy.Subscriber('/image_color',
Image,
self.image_cb,
queue_size=1,
tcp_nodelay=True)
rospy.spin()
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, waypoints):
self.waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in waypoints.waypoints]
self.waypoints_tree = KDTree(self.waypoints_2d)
def traffic_cb(self, msg):
self.lights = msg.lights
def image_cb(self, msg):
"""Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.has_image = True
self.camera_image = msg
light_wp, state = self.process_traffic_lights()
'''
Publish upcoming red lights at camera frequency.
Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
of times till we start using it. Otherwise the previous stable state is
used.
'''
if self.state != state:
self.state_count = 0
self.state = state
elif self.state_count >= STATE_COUNT_THRESHOLD:
self.last_state = self.state
light_wp = light_wp if state == TrafficLight.RED else -1
self.last_wp = light_wp
self.upcoming_red_light_pub.publish(Int32(light_wp))
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
self.state_count += 1
self.has_image = False
'''
def get_closest_waypoint(self, pose):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
# TODO implement
closest_idx = self.waypoints_tree.query([pose.position.x, pose.position.y], 1)[1]
return closest_idx
'''
def get_closest_waypoint(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
x (float): position to match a waypoint to
y (float): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
# TODO implement
closest_idx = self.waypoints_tree.query([x, y], 1)[1]
return closest_idx
def get_light_state(self, light):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
#return light.state
if not self.has_image:
self.prev_light_loc = None
return TrafficLight.UNKNOWN
if self.config['is_site']:
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "rgb8")
else:
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "rgb8")
# Get classification
result = self.light_classifier.get_classification(cv_image)
return result
def process_traffic_lights(self):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
closest_light = None
line_wp_idx = None
# List of positions that correspond to the line to stop in front of for a given intersection
stop_line_positions = self.config['stop_line_positions']
light_state_to_str = {0: "RED", 2: "GREEN", 1: "YELLOW", 4: "UNKNOWN"}
# TODO: just for testing, remove it
#if self.config['is_site']:
# state = self.get_light_state(closest_light)
# rospy.logwarn("Nearest traffic light state: {0}".format(light_state_to_str[int(state)]))
# TODO: end of testing
if self.pose and self.waypoints and self.waypoints_tree:
car_wp_idx = self.get_closest_waypoint(self.pose.pose.position.x,
self.pose.pose.position.y)
# TODO find the closest visible traffic light (if one exists)
diff = len(self.waypoints.waypoints)
for i, light in enumerate(self.lights):
# Get stop line waypoint index
line = stop_line_positions[i]
temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
# Find closest stop line waypoint index
d = temp_wp_idx - car_wp_idx
if 0 <= d < diff:
diff = d
closest_light = light
line_wp_idx = temp_wp_idx
if closest_light:
state = self.get_light_state(closest_light)
rospy.logwarn("Nearest traffic light state: {0}".format(
light_state_to_str[state]))
return line_wp_idx, state
#self.waypoints = None
rospy.logwarn("Nearest traffic light state: UNKNOWN")
return -1, TrafficLight.UNKNOWN
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.camera_image = None
self.lights = []
self.waypoints_2d = None
self.waypoints_tree = None
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
self.is_sim = not self.config["is_site"]
self.is_sim = True
print("Is Sim %d" % self.is_sim)
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.bridge = CvBridge()
self.light_classifier = TLClassifier(self.is_sim)
self.listener = tf.TransformListener()
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
# rospy.spin()
self.ros_spin()
def ros_spin(self):
rate = rospy.Rate(100)
while not rospy.is_shutdown():
'''Publish upcoming red lights at camera frequency.
Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
of times till we start using it. Otherwise the previous stable state is
used.
'''
if self.pose is not None and self.waypoints is not None and self.camera_image is not None:
light_wp, state = self.process_traffic_lights()
# print("Light waypoint Index: ",light_wp, "Traffic Light: ", TrafficLight.RED, state)
if self.state != state:
self.state_count = 0
self.state = state
elif self.state_count >= STATE_COUNT_THRESHOLD:
self.last_state = self.state
light_wp = light_wp if state == TrafficLight.RED else -1
self.last_wp = light_wp
self.upcoming_red_light_pub.publish(Int32(light_wp))
# print(light_wp,self.state_count)
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
self.state_count += 1
rate.sleep()
def pose_cb(self, msg):
self.pose = msg
# Extracts the x and y coordinates of a waypoint
def waypoint_xy(self, waypoint):
position = waypoint.pose.pose.position
return [position.x, position.y]
def waypoints_cb(self, waypoints):
self.waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [
self.waypoint_xy(waypoint) for waypoint in waypoints.waypoints
]
self.waypoints_tree = KDTree(self.waypoints_2d)
def traffic_cb(self, msg):
self.lights = msg.lights
def image_cb(self, msg):
"""Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.has_image = True
self.camera_image = msg
print('get camera image', rospy.Time.now())
# Commented out to test for lag
# light_wp, state = self.process_traffic_lights()
def get_closest_waypoint(self, pose_x, pose_y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args: pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
# Use KDTree to search through waypoints
closest_idx = self.waypoints_tree.query([pose_x, pose_y], 1)[1]
return closest_idx
def get_light_state(self, light):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
if not self.has_image:
self.prev_light_loc = None
return False
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
# Get classification
classification_result = self.light_classifier.get_classification(
cv_image)
rospy.loginfo("ground truth: %s, classifier result: %s", light.state,
classification_result)
return classification_result
# return light.state
def process_traffic_lights(self):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
closest_light = None
light_wp_idx = None
# List of positions that correspond to the line to stop in front of for a given intersection
stop_line_positions = self.config['stop_line_positions']
if (self.pose):
car_wp_idx = self.get_closest_waypoint(self.pose.pose.position.x,
self.pose.pose.position.y)
# print(car_wp_idx)
# TODO find the closest visible traffic light (if one exists)
diff = len(self.waypoints.waypoints)
# print(self.pose.pose.position.x,self.pose.pose.position.y, car_wp_idx)
for i, light in enumerate(self.lights):
# Get stop line waypoint index
line = stop_line_positions[i]
# print(line[0], line[1])
temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
# Find Closest stop line waypoint index
d = temp_wp_idx - car_wp_idx
if d >= 0 and d < diff:
diff = d
closest_light = light
light_wp_idx = temp_wp_idx
if closest_light:
state = self.get_light_state(closest_light)
# print("Light waypoint Index: ",light_wp_idx, "Traffic Light: ",TrafficLight.RED, state)
return light_wp_idx, state
# self.waypoints = None
return -1, TrafficLight.UNKNOWN
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.base_waypoints = None
self.waypoints_2d = None
self.waypoint_tree = None
self.camera_image = None
self.lights = []
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.bridge = CvBridge()
self.light_classifier = TLClassifier()
self.listener = tf.TransformListener()
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
rospy.spin()
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, waypoints):
if not self.waypoint_tree:
self.base_waypoints = waypoints
self.waypoints_2d = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in waypoints.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d)
def traffic_cb(self, msg):
self.lights = msg.lights
def image_cb(self, msg):
"""Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.has_image = True
self.camera_image = msg
light_wp_idx, state = self.process_traffic_lights()
'''
Publish upcoming red lights at camera frequency.
Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
of times till we start using it. Otherwise the previous stable state is
used.
'''
if self.state != state:
self.state_count = 0
self.state = state
elif self.state_count >= STATE_COUNT_THRESHOLD:
self.last_state = self.state
light_wp_idx = light_wp_idx if state == TrafficLight.RED else -1
self.last_wp = light_wp_idx
self.upcoming_red_light_pub.publish(Int32(light_wp_idx))
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
self.state_count += 1
def get_closest_waypoint_idx(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.base_waypoints
"""
if self.waypoint_tree is None:
return -1
else:
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
return closest_idx
def get_light_state(self, light):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
# For testing, just return the light state
# TODO: remove the next line after testing
# return light.state
if (not self.has_image):
self.prev_light_loc = None
return False
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
#Get classification
prob = self.calc_probability_of_traffic_light_in_camera_view(light)
return self.light_classifier.get_classification(cv_image, prob)
def calc_probability_of_traffic_light_in_camera_view(self, light):
light_pos = light.pose.pose.position
car_pos = self.pose.pose.position
car_roll, car_pitch, car_yaw = euler_from_quaternion([
self.pose.pose.orientation.x, self.pose.pose.orientation.y,
self.pose.pose.orientation.z, self.pose.pose.orientation.w
])
car_w = car_yaw
car_direction = np.array([np.cos(car_w), np.sin(car_w)])
light_direction = np.array(
[light_pos.x - car_pos.x, light_pos.y - car_pos.y])
cos = np.dot(car_direction, light_direction) / np.linalg.norm(
car_direction) / np.linalg.norm(light_direction)
angle = np.arccos(np.clip(cos, -1, 1))
dx = light_pos.x - car_pos.x
dy = light_pos.y - car_pos.y
dist_to_light = np.sqrt(dx**2 + dy**2)
detect_dist = dist_to_light - 10 # m - if the car is below the traffic light, it will not see the light
detect_angle = abs(angle)
prob_dist = (200 - detect_dist) / 200
prob_angle = (np.radians(45) - detect_angle) / np.radians(45)
probability = min(max(prob_dist * prob_angle, 0),
1) if detect_angle > 0. and detect_dist > 0. else 0.
# rospy.loginfo("[TLD] probability of traffic light {0} ({1} * {2})".format(probability, prob_dist, prob_angle))
return probability
def process_traffic_lights(self):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
closest_light = None
closest_line_wp_idx = None
# List of positions that correspond to the line to stop in front of for a given intersection
stop_line_positions = self.config['stop_line_positions']
if not None in (self.pose, self.waypoint_tree):
car_wp_idx = self.get_closest_waypoint_idx(
self.pose.pose.position.x, self.pose.pose.position.y)
#TODO find the closest visible traffic light (if one exists)
min_diff = len(self.base_waypoints.waypoints)
for i, light in enumerate(self.lights):
# Get stop line waypoint index
line = stop_line_positions[i]
line_wp_idx = self.get_closest_waypoint_idx(line[0], line[1])
# Find closest stop line waypoint index
d = line_wp_idx - car_wp_idx
if d >= 0 and d < min_diff:
min_diff = d
closest_light = light
closest_line_wp_idx = line_wp_idx
else:
rospy.logwarn("[TLD] pose or waypoints is missing")
if closest_light:
state = self.get_light_state(closest_light)
return closest_line_wp_idx, state
else:
rospy.logwarn("[TLD] no closest light found")
return -1, TrafficLight.UNKNOWN
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.camera_image = None
self.waypoints_2d = None
self.waypoint_tree = None
self.has_image = False
self.lights = []
self.last_light_state = TrafficLight.UNKNOWN
self.immediate_light_state = TrafficLight.UNKNOWN
self.immediate_light_state_count = 0
self.light_classifier_throttle_time = rospy.get_rostime()
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray, self.traffic_cb)
sub6 = rospy.Subscriber('/image_color', Image, self.image_cb, queue_size=1)
self.test_pub = rospy.Publisher('/dummy_image', Image, queue_size=1)
self.light_pub = rospy.Publisher('/tl_filtered', String, queue_size=1)
config_string = rospy.get_param("/traffic_light_config")
# config also contains an "is_site" bool variable to check if the car is running in simulator or site
self.config = yaml.load(config_string)
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint', Int32, queue_size=1)
self.bridge = CvBridge()
self.light_classifier = TLClassifier()
self.light_classifier.find_objects(np.zeros((2,2,3), np.uint8), 0.5, [10.]) # Hit the classifier to get TF through its init
self.listener = tf.TransformListener()
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
self.main_loop()
#rospy.spin()
def main_loop(self):
while not rospy.is_shutdown():
if(self.has_image):
light_wp, state = self.process_traffic_lights()
if self.state != state:
self.state_count = 0
self.state = state
elif self.state_count >= STATE_COUNT_THRESHOLD:
self.last_state = self.state
light_wp = light_wp if state == TrafficLight.RED else -1
self.last_wp = light_wp
self.upcoming_red_light_pub.publish(Int32(light_wp))
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
self.state_count += 1
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, waypoints):
self.waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[waypoint.pose.pose.position.x, waypoint.pose.pose.position.y] for waypoint in waypoints.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d) # store in a tree structure for easier access
def traffic_cb(self, msg):
self.lights = msg.lights
def image_cb(self, msg):
"""Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.has_image = True
self.camera_image = msg
self.test_pub.publish(self.camera_image)
def get_closest_waypoint(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
# find the x, y position in the waypoints in the KDTree
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
return closest_idx
def get_light_state(self, light):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
# just for testing
#return light.state
if(not self.has_image):
self.prev_light_loc = None
return False
#Get classification
# Throttle the incoming requests
#if self.light_classifier_throttle_time + rospy.Duration(0.05) < rospy.get_rostime():
#cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
#self.last_light_state = self.light_classifier.get_classification(cv_image)
#self.light_classifier_throttle_time = rospy.get_rostime()
#return self.last_light_state
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
self.immediate_light_state = self.light_classifier.get_classification(cv_image)
self.light_classifier_throttle_time = rospy.get_rostime()
# Filter out noise from the classifier
if (self.immediate_light_state != self.last_light_state):
self.immediate_light_state_count += 1
if (self.immediate_light_state_count > 4):
#if (self.last_light_state == 1 and self.immediate_light_state != 0):
# self.last_light_state = self.last_light_state
#else:
self.last_light_state = self.immediate_light_state
self.immediate_light_state_count = 0
self.light_pub.publish(str(self.last_light_state))
return self.last_light_state
#else:
# return self.last_light_state
def process_traffic_lights(self):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
closest_light = None
line_wp_idx = None
# List of positions that correspond to the line to stop in front of for a given intersection
stop_line_positions = self.config['stop_line_positions']
if(self.pose):
x = self.pose.pose.position.x
y = self.pose.pose.position.y
car_wp_idx = self.get_closest_waypoint(x, y)
#TODO find the closest visible traffic light (if one exists)
diff = len(self.waypoints.waypoints)
# iterate through the traffic lights to find the closest
for i, light in enumerate(self.lights):
# Get stop line waypoint index
line = stop_line_positions[i]
temp_wp_idx = self.get_closest_waypoint(line[0], line[1])
# Find closest stop line waypoint idx
d = temp_wp_idx - car_wp_idx
# linear check to find the closest
if d >= 0 and d < diff:
diff = d
closest_light = light
line_wp_idx = temp_wp_idx
if closest_light:
state = self.get_light_state(closest_light)
return line_wp_idx, state
return -1, TrafficLight.UNKNOWN
X, Y, THETA = 0, 0.2, 0
# 速度，车轮间距，间隔时间
V, L, DT = 0.15, 0.4, 0.1
# 初始化质点模型
robot = particle.MODEL(X, Y, THETA, V, L, DT)
# 配置增量式PID控制器参数
pid = incremental_pid.CU(0.8, 0.05, 20)
# 初始化设备当前位置量
robot_state = np.zeros(2)
for _ in range(COUNT):
# 获取当前位置
robot_state[0], robot_state[1] = robot.x, robot.y
_, ind = refer_tree.query(robot_state)
# 计算当前与路径的横向距离偏差
_, dy = refer_path[ind] - robot_state
# 计算下一状态与当前路径x方向的期望距离
ndx, ndy = refer_path[ind + 1] - refer_path[ind]
dx = V * np.cos(math.atan2(ndy, ndx)) * robot.dt
# 计算期望偏转角
alpha = math.atan2(dy, dx)
# 计算偏转角偏差
e = alpha - robot.theta
# 更新偏差
class Simulation(object):
def __init__(self, path, mag_zeropoint, exposure):
"""
Parses a Stuff input file. See https://www.astromatic.net/software/stuff
:param path:
Path to the .list simulation data
:param mag_zeropoint:
Used to compute the flux
:param exposure:
Used to compute the flux
"""
logger.debug(f'Loading stuff list from {path}')
stars_raw, galaxies_raw = [], []
with open(path, 'r') as lst:
for entry in map(str.split, map(str.strip, lst.readlines())):
if entry[0] == '100':
stars_raw.append(entry)
elif entry[0] == '200':
galaxies_raw.append(entry)
else:
raise Exception(f'Unexpected type code {entry[0]}')
logger.debug(f'Loaded {len(stars_raw)} stars')
logger.debug(f'Loaded {len(galaxies_raw)} galaxies')
self.__stars = np.recarray((len(stars_raw), ),
dtype=[('ra', float), ('dec', float),
('mag', float), ('flux', float)])
self.__galaxies = np.recarray((len(galaxies_raw), ),
dtype=[('ra', float), ('dec', float),
('mag', float), ('flux', float),
('bt_ratio', float),
('bulge', float),
('bulge_aspect', float),
('bulge_angle', float),
('disk', float),
('disk_aspect', float),
('disk_angle', float),
('redshift', float),
('type', float)])
for i, s in enumerate(stars_raw):
self.__stars[i].ra, self.__stars[i].dec, self.__stars[
i].mag = float(s[1]), float(s[2]), float(s[3])
self.__stars.flux = exposure * np.power(
10, (self.__stars.mag - mag_zeropoint) / -2.5)
for i, g in enumerate(galaxies_raw):
self.__galaxies[i].ra, self.__galaxies[i].dec = float(g[1]), float(
g[2])
self.__galaxies[i].mag = float(g[3])
self.__galaxies[i].bt_ratio = float(g[4])
self.__galaxies[i].bulge, self.__galaxies[
i].bulge_aspect, self.__galaxies[i].bulge_angle = float(
g[5]), float(g[6]), float(g[7])
self.__galaxies[i].disk, self.__galaxies[
i].disk_aspect, self.__galaxies[i].disk_angle = float(
g[8]), float(g[9]), float(g[10])
self.__galaxies[i].redshift = float(g[11])
self.__galaxies[i].type = float(g[12])
self.__galaxies.flux = exposure * np.power(
10, (self.__galaxies.mag - mag_zeropoint) / -2.5)
all_coords = np.column_stack([
np.append(self.__stars.ra, self.__galaxies.ra),
np.append(self.__stars.dec, self.__galaxies.dec)
])
self.__kdtree = KDTree(all_coords)
self.__all_mags = np.append(self.__stars.mag, self.__galaxies.mag)
@property
def stars(self):
return self.__stars
@property
def galaxies(self):
return self.__galaxies
@property
def magnitudes(self):
"""
:return: All the magnitudes on the simulation: first stars, then galaxies.
"""
return self.__all_mags
def get_star_count(self):
"""
:return: Number of stars on the simulation.
"""
return len(self.__stars)
def get_galaxy_count(self):
"""
:return: Number of galaxies on the simulation.
"""
return len(self.__galaxies)
def get_closest(self, alpha, delta):
"""
Find the closest source to the catalog entries. This can be used to cross-relate
the entries from the Stuff simulation and a catalog
:param alpha:
Alpha coordinates
:param delta:
Delta coordinates
:return:
A numpy array with the columns:
* distance: distance to closest source
* catalog: corresponding index on the catalog
* source: corresponding index on the simulation list
* magnitude: magnitude of the closest source
* is_star: True if the closest source is a star
* is_galaxy: True if the closest source is a galaxy
"""
assert len(alpha) == len(delta)
closest = np.recarray((len(alpha), ),
dtype=[('distance', float), ('catalog', int),
('source', int), ('magnitude', float),
('is_star', bool), ('is_galaxy', bool)])
closest.catalog[:] = np.arange(0, len(alpha))
closest.distance[:], closest.source[:] = self.__kdtree.query(
np.column_stack([alpha, delta]))
closest.magnitude[:] = self.__all_mags[closest.source]
closest.is_star[:] = closest.source < self.get_star_count()
closest.is_galaxy[:] = np.logical_not(closest.is_star[:])
return closest
def fast_knn(data,
k=3,
eps=0,
p=2,
distance_upper_bound=np.inf,
leafsize=10,
idw_fn=idw.shepards,
init_impute_fn=mean):
""" Impute using a variant of the nearest neighbours approach
Basic idea: Impute array with a passed in initial impute fn (mean impute)
and then use the resulting complete array to construct a KDTree. Use this
KDTree to compute nearest neighbours. After finding `k` nearest
neighbours, take the weighted average of them. Basically, find the nearest
row in terms of distance
This approach is much, much faster than the other implementation (fit+transform
for each subset) which is almost prohibitively expensive.
Parameters
----------
data: numpy.ndarray
2D matrix to impute.
k: int, optional
Parameter used for method querying the KDTree class object. Number of
neighbours used in the KNN query. Refer to the docs for
[`scipy.spatial.KDTree.query`]
(https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.query.html).
eps: nonnegative float, optional
Parameter used for method querying the KDTree class object. From the
SciPy docs: "Return approximate nearest neighbors; the kth returned
value is guaranteed to be no further than (1+eps) times the distance to
the real kth nearest neighbor". Refer to the docs for
[`scipy.spatial.KDTree.query`]
(https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.query.html).
p : float, 1<=p<=infinity, optional
Parameter used for method querying the KDTree class object. Straight from the
SciPy docs: "Which Minkowski p-norm to use. 1 is the
sum-of-absolute-values Manhattan distance 2 is the usual Euclidean
distance infinity is the maximum-coordinate-difference distance". Refer to
the docs for
[`scipy.spatial.KDTree.query`]
(https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.query.html).
distance_upper_bound : nonnegative float, optional
Parameter used for method querying the KDTree class object. Straight
from the SciPy docs: "Return only neighbors within this distance. This
is used to prune tree searches, so if you are doing a series of
nearest-neighbor queries, it may help to supply the distance to the
nearest neighbor of the most recent point." Refer to the docs for
[`scipy.spatial.KDTree.query`]
(https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.KDTree.query.html).
leafsize: int, optional
Parameter used for construction of the `KDTree` class object. Straight from
the SciPy docs: "The number of points at which the algorithm switches
over to brute-force. Has to be positive". Refer to the docs for
[`scipy.spatial.KDTree`](https://docs.scipy.org/doc/scipy-0.14.0/reference/generated/scipy.spatial.KDTree.html)
for more information.
idw_fn: fn, optional
Function that takes one argument, a list of distances, and returns weighted percentages. You can define a custom
one or bootstrap from functions defined in `impy.util.inverse_distance_weighting` which can be using
functools.partial, for example: `functools.partial(impy.util.inverse_distance_weighting.shepards, power=1)`
init_impute_fn: fn, optional
Returns
-------
numpy.ndarray
Imputed data.
Examples
--------
>>> data = np.arange(25).reshape((5, 5)).astype(np.float)
>>> data[0][2] = np.nan
>>> data
array([[ 0., 1., nan, 3., 4.],
[ 5., 6., 7., 8., 9.],
[10., 11., 12., 13., 14.],
[15., 16., 17., 18., 19.],
[20., 21., 22., 23., 24.]])
>> fast_knn(data, k=1) # Weighted average (by distance) of nearest 1 neighbour
array([[ 0., 1., 7., 3., 4.],
[ 5., 6., 7., 8., 9.],
[10., 11., 12., 13., 14.],
[15., 16., 17., 18., 19.],
[20., 21., 22., 23., 24.]])
>> fast_knn(data, k=2) # Weighted average of nearest 2 neighbours
array([[ 0. , 1. , 10.08608891, 3. , 4. ],
[ 5. , 6. , 7. , 8. , 9. ],
[10. , 11. , 12. , 13. , 14. ],
[15. , 16. , 17. , 18. , 19. ],
[20. , 21. , 22. , 23. , 24. ]])
>> fast_knn(data, k=3)
array([[ 0. , 1. , 13.40249283, 3. , 4. ],
[ 5. , 6. , 7. , 8. , 9. ],
[10. , 11. , 12. , 13. , 14. ],
[15. , 16. , 17. , 18. , 19. ],
[20. , 21. , 22. , 23. , 24. ]])
>> fast_knn(data, k=5) # There are at most only 4 neighbours. Raises error
...
IndexError: index 5 is out of bounds for axis 0 with size 5
"""
nan_xy = matrix.nan_indices(data)
data_c = init_impute_fn(data)
kdtree = KDTree(data_c, leafsize=leafsize)
for x_i, y_i in nan_xy:
distances, indices = kdtree.query(
data_c[x_i],
k=k + 1,
eps=eps,
p=p,
distance_upper_bound=distance_upper_bound)
# Will always return itself in the first index. Delete it.
distances, indices = distances[1:], indices[1:]
# Add small constant to distances to avoid division by 0
distances += 1e-3
weights = idw_fn(distances)
# Assign missing value the weighted average of `k` nearest neighbours
data[x_i][y_i] = np.dot(weights, [data_c[ind][y_i] for ind in indices])
return data
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.base_waypoints = None
self.waypoints_2d = None
self.waypoint_tree = None
#the location of stopline
self.stopline_2d = [[1148.56, 1184.65], [1559.2, 1158.43],
[2122.14, 1526.79], [2175.237, 1795.71],
[1493.29, 2947.67], [821.96, 2905.8],
[161.76, 2303.82], [351.84, 1574.65]]
self.stopline_tree = KDTree(self.stopline_2d)
#set the light's initial state not red
self.red_light = [-1, -1, -1, -1, -1, -1, -1, -1]
self.lights = []
rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.loop()
def loop(self):
rate = rospy.Rate(50)
while not rospy.is_shutdown():
if self.pose and self.base_waypoints:
self.publish_redlight_idx()
rate.sleep()
def publish_redlight_idx(self):
#get the closest traffic light's index in stopline_2d
closest_light_idx = self.get_closest_redlight_idx()
#get the state of closest light, if it's red light, get it's index in base_waypoint and publish, if not red light publish -1.
redlight_state = self.red_light[closest_light_idx]
if redlight_state == 1 and self.waypoint_tree:
closest_red_light_idx = self.get_closest_waypoint_idx(
closest_light_idx)
self.upcoming_red_light_pub.publish(Int32(closest_red_light_idx))
else:
self.upcoming_red_light_pub.publish(Int32(-1))
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, waypoints):
self.base_waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in waypoints.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d)
def traffic_cb(self, msg):
self.lights = msg.lights
light_state = [trafficlight.state for trafficlight in msg.lights]
for i, state in enumerate(light_state):
if state == 0:
self.red_light[i] = 1
else:
self.red_light[i] = -1
def get_closest_redlight_idx(self):
#use this function to get the closest traffic light's index in stopline_2d
x = self.pose.pose.position.x
y = self.pose.pose.position.y
closest_idx = self.stopline_tree.query([x, y], 1)[1]
# check if closest is ahead or behind car
closest_location = self.stopline_2d[closest_idx]
pre_closest_location = self.stopline_2d[closest_idx - 1]
cl_vect = np.array(closest_location)
pre_vect = np.array(pre_closest_location)
pos_vect = np.array([x, y])
if np.dot(cl_vect - pre_vect, pos_vect - cl_vect) > 0:
closest_idx = (closest_idx + 1) % len(self.stopline_2d)
return closest_idx
def get_closest_waypoint_idx(self, closest_light_idx):
#use this function to get the redlight's index in waypoints_2d
x = self.stopline_2d[closest_light_idx][0]
y = self.stopline_2d[closest_light_idx][1]
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
# check if closest is ahead or behind car
closest_waypoint = self.waypoints_2d[closest_idx]
pre_closest_waypoint = self.waypoints_2d[closest_idx - 1]
cl_vect = np.array(closest_waypoint)
pre_vect = np.array(pre_closest_waypoint)
pos_vect = np.array([x, y])
if np.dot(cl_vect - pre_vect, pos_vect - cl_vect) > 0:
closest_idx = (closest_idx + 1) % len(self.waypoints_2d)
return closest_idx
class WaypointUpdater(object):
def __init__(self):
rospy.init_node('waypoint_updater')
rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
rospy.Subscriber('/traffic_waypoint', Int32, self.traffic_cb)
self.final_waypoints_pub = rospy.Publisher('final_waypoints', Lane, queue_size=1)
self.pose = None
self.stopline_wp_idx = -1
self.base_waypoints = None
self.waypoints_2d = None
self.waypoints_tree = None
self.loop()
def loop(self):
rate = rospy.Rate(50)
while not rospy.is_shutdown():
if self.pose and self.base_waypoints:
self.publish_waypoints()
rate.sleep()
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, waypoints):
self.base_waypoints = waypoints
if not self.waypoints_2d:
self.waypoints_2d = [[waypoint.pose.pose.position.x, waypoint.pose.pose.position.y] for waypoint in waypoints.waypoints]
self.waypoints_tree = KDTree(self.waypoints_2d)
def get_closest_waypoint_idx(self):
x = self.pose.pose.position.x
y = self.pose.pose.position.y
closest_idx = self.waypoints_tree.query([x,y],1)[1]
closest_coord = self.waypoints_2d[closest_idx]
prev_coord = self.waypoints_2d[closest_idx-1]
cl_vec = np.array(closest_coord)
prev_vec = np.array(prev_coord)
pos_vec = np.array([x,y])
val = np.dot(cl_vec-prev_vec, pos_vec-cl_vec)
if val > 0:
closest_idx = (closest_idx+1) % len(self.waypoints_2d)
return closest_idx
def traffic_cb(self, msg):
self.stopline_wp_idx = msg.data
pass
def obstacle_cb(self, msg):
# TODO: Callback for /obstacle_waypoint message. We will implement it later
pass
def get_waypoint_velocity(self, waypoint):
return waypoint.twist.twist.linear.x
def set_waypoint_velocity(self, waypoints, waypoint, velocity):
waypoints[waypoint].twist.twist.linear.x = velocity
def publish_waypoints(self):
final_lane = self.generate_lane()
self.final_waypoints_pub.publish(final_lane)
def generate_lane(self):
lane = Lane()
closest_idx = self.get_closest_waypoint_idx()
base_waypoints = self.base_waypoints.waypoints[closest_idx:closest_idx+LOOKAHEAD_WPS]
if self.stopline_wp_idx == -1 or (self.stopline_wp_idx >= (closest_idx+LOOKAHEAD_WPS)):
lane.waypoints = base_waypoints
else:
lane.waypoints = self.decelerate_waypoints(base_waypoints, closest_idx)
return lane
def decelerate_waypoints(self, waypoints, closest_idx):
new_wp = []
for i, wp in enumerate(waypoints):
p = Waypoint()
p.pose = wp.pose
stop_idx = max(self.stopline_wp_idx-closest_idx-2, 0)
dist = self.distance(waypoints, i, stop_idx)
vel = math.sqrt(2*MAX_DECEL*dist)
if vel < 1.0:
vel = 0.
p.twist.twist.linear.x = min(vel, wp.twist.twist.linear.x)
new_wp.append(p)
return new_wp
def distance(self, waypoints, wp1, wp2):
dist = 0
dl = lambda a, b: math.sqrt((a.x-b.x)**2 + (a.y-b.y)**2 + (a.z-b.z)**2)
for i in range(wp1, wp2+1):
dist += dl(waypoints[wp1].pose.pose.position, waypoints[i].pose.pose.position)
wp1 = i
return dist
MPIrank,
time.time() - tb))
if args.testMode:
print('TEST MODE')
sys.exit(1)
###########################################################
### GrC PFs to GoC AD for glutamatergic synapses ###
###########################################################
vprint('brep.py rank {0} PFs to GoCs'.format(MPIrank))
tb = time.time()
tree = KDTree(PFs[:, 2:])
K = min(len(PFs), K4T)
distKDTree, idxKDTree = tree.query(GoCadend, k=K)
results = np.stack([
np.repeat(np.arange(len(GoCadend)), K).reshape(len(GoCadend), K),
PFs[idxKDTree, 1],
distKDTree,
idxKDTree,
]).transpose(1, 2, 0)
results = results.reshape(results.shape[0] * results.shape[1],
results.shape[2])
results = results[results[:, 2] <= h.PFtoGoCzone]
results = results[np.unique(results[:, :2], axis=0, return_index=True)[1]]
if hasattr(h, "BREP_dendcoeff") and h.BREP_dendcoeff > 0:
dendLen = h.BREP_dendcoeff * np.linalg.norm(GoCadend[results[:, 0].astype(
np.int64)] - GoC[GoCadendIdx[results[:, 0].astype(np.int64)]],
axis=1)
def plot_fig(output, dx, E, B, cutoff, energies, alpha0, text, xlim, ylim):
fig = plot.figure(figsize=(244.0 / 72, 135.0 / 72))
ax = fig.add_subplot(1, 1, 1, aspect='equal')
ax.set_xlim(*xlim)
ax.set_ylim(*ylim)
x0 = np.zeros(3)
for KE in energies:
for dalpha, linewidth, opacity in zip(np.array(
(-5 * deg, 0, 5 * deg)), [0.2, 0.4, 0.2], [0.5, 1.0, 0.5]):
alpha = alpha0 + dalpha
n0 = np.array((-np.sin(alpha), np.cos(alpha), 0))
traj = calc_electron_trajectory(n0, x0, KE, E, B, dx, cutoff)
ax.plot(traj[0] / mm,
traj[1] / mm,
linewidth=linewidth,
color='#0083b8',
alpha=opacity)
alpha = alpha0 - 5 * deg
n0 = np.array((-np.sin(alpha), np.cos(alpha), 0))
x1 = calc_electron_trajectory(n0, x0, energies[-1], E, B, 0.5 * mm,
cutoff)[0:2].T
N = len(x1)
x1 = x1[N // 4:N - N // 4]
alpha = alpha0 + 5 * deg
n0 = np.array((-np.sin(alpha), np.cos(alpha), 0))
x2 = calc_electron_trajectory(n0, x0, energies[-1], E, B, 0.5 * mm,
cutoff)[0:2].T
N = len(x2)
x2 = x2[N // 4:N - N // 4]
tree = KDTree(x1)
x2_scan = np.array([tree.query(q) for q in x2])
i2 = np.argmin(x2_scan.T[0])
dist, i1 = tree.query(x2[i2])
x1_f = 1.05 * x1[i1]
line_length = norm(x1_f)
x1_0 = 0.6 * line_length * np.array((-np.sin(alpha0), np.cos(alpha0)))
ax.plot([0, x1_f[0] / mm], [0, x1_f[1] / mm],
color='#888888',
linewidth=0.4,
zorder=-10)
ax.text(x1_f[0] / mm,
x1_f[1] / mm,
r' $\alpha_f = {:.1f}^\circ$'.format(90 -
np.arctan(x1_f[1] / x1_f[0]) /
deg),
fontsize=7,
verticalalignment='center')
ax.plot([0, x1_0[0] / mm], [0, x1_0[1] / mm],
color='#888888',
linewidth=0.4,
zorder=-10)
alpha0_label = r' $\alpha_0 = {:.1f}^\circ$'.format(alpha0 / deg)
ax.text(x1_0[0] / mm,
x1_0[1] / mm,
alpha0_label,
fontsize=7,
verticalalignment='center',
horizontalalignment='right')
plot.xlabel(r'$x$ (mm)', labelpad=0.0)
plot.ylabel(r'$z$ (mm)', labelpad=0.0)
ax.text(0.03,
0.96,
text,
fontsize=7,
verticalalignment='top',
transform=ax.transAxes)
ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
output.savefig(fig, transparent=True)
def scatter_to_regulargrid(xcoords=None,
ycoords=None,
ncellx=None,
ncelly=None,
values=None,
method='nearest',
maskland_dist=None):
"""
interpolates scatter values (x,y,z) or meshgrids to regular grid
Parameters
----------
xcoords : TYPE, optional
DESCRIPTION. The default is None.
ycoords : TYPE, optional
DESCRIPTION. The default is None.
ncellx : TYPE, optional
DESCRIPTION. The default is None.
ncelly : TYPE, optional
DESCRIPTION. The default is None.
values : TYPE, optional
DESCRIPTION. The default is None.
method : TYPE, optional
DESCRIPTION. The default is 'nearest'.
maskland_dist : TYPE, optional
Distance used to mask out land cells. Value differs per model resolution. For sperical use maskland_dist of eg 0.01, for cartesian use maskland_dist of eg 100. The default is None.
Returns
-------
x_grid : TYPE
DESCRIPTION.
y_grid : TYPE
DESCRIPTION.
value_grid : TYPE
DESCRIPTION.
"""
import numpy as np
from scipy.interpolate import griddata
from scipy.spatial import KDTree
reg_x_vec = np.linspace(np.min(xcoords), np.max(xcoords), ncellx)
reg_y_vec = np.linspace(np.min(ycoords), np.max(ycoords), ncelly)
x_grid, y_grid = np.meshgrid(reg_x_vec, reg_y_vec)
#first replace masked value with nan (mask is not used in griddata)
try:
values[values.mask] = np.nan
except:
pass
value_grid = griddata((xcoords, ycoords),
values, (x_grid, y_grid),
method=method)
if maskland_dist is not None:
#Mask out land values with maskland_dist
#https://stackoverflow.com/questions/10456143/get-only-valid-points-in-2d-interpolation-of-cloud-point-using-scipy-numpy
tree = KDTree(np.c_[xcoords, ycoords])
dist, _ = tree.query(np.c_[x_grid.ravel(), y_grid.ravel()], k=1)
dist = dist.reshape(x_grid.shape)
value_grid[dist > maskland_dist] = np.nan
return x_grid, y_grid, value_grid
class kNNBatchDataset(Dataset):
def __init__(self, *args, **kwargs):
super(kNNBatchDataset, self).__init__(*args, **kwargs)
self.tree = None
self.buildKD()
self.center_idx = 0
def buildKD(self):
logging.info(" -- Building kD-Tree with %d points..." % len(self))
self.tree = KDTree(self._xyz[:, :2], leafsize=100) # build only on x/y
logging.info(" --- kD-Tree built.")
def getBatches(self, batch_size=1):
centers = []
for i in range(batch_size):
if self.currIdx >= self.num_batches:
break
centers.append([
self.xmin + self.spacing / 2 +
(self.currIdx // self.num_rows) * self.spacing,
self.ymin + self.spacing / 2 +
(self.currIdx % self.num_rows) * self.spacing
])
self.currIdx += 1
print(centers)
if centers:
_, idx = self.tree.query(centers, k=self.k)
return self.points_and_features[idx, :], self.labels[idx]
else:
return None, None
def getBatches_Point(self, batch_size=1, num_point=1024, num_grid=32):
batch_points = []
batch_labels = []
batch_labels = torch.zeros([batch_size, 1],
dtype=torch.float) # batch size * 1
batch_voxels = []
batch_points = torch.zeros([batch_size, 3, num_point],
dtype=torch.float)
for i in range(batch_size):
points = []
voxels = []
tmp_index = self.center_idx
if (self.shuffle == True):
tmp_index = self.index_for_train[self.center_idx]
_, idx = self.tree.query(self.points_and_features[tmp_index, :2],
k=num_point)
label = self.labels[tmp_index]
# print(label)
batch_labels[i] = torch.from_numpy(
np.array(label).astype(np.float32))
point_knn = np.full((num_point, 3), 1)
point_knn[:, :] = self.points_and_features[idx, :3]
batch_points[i, :, :] = torch.from_numpy(
point_knn.astype(np.float32).T)
self.center_idx += 1
# for j in range(len(num_point)):
# point_knn = np.full((num_point[j], 3), 1)
# point_knn[:, :] = self.points_and_features[idx[:num_point[j]], :3]
# # point_knn = torch.from_numpy(np.array(point_knn).astype(np.float32))
# # print(point_knn.shape)
# batch_voxels[j, i, :, :, :, :] = torch.from_numpy(np.array(voxelization(point_knn,
# vox_size=num_grid)).astype(np.float32)) # resolution * batchsize * ch * grid * grid * grid
# points.append(point_knn)
# print(point_knn)
# print(point_knn.shape)
# print(batch_voxels.shape)
# batch_points.append(points) # batchsize * resolution * num.of points * xyz
# batch_points = torch.from_numpy(np.array(batch_points).astype(np.float32)) # batchsize * resolution * num.of points * xyz
# for i in range(batch_size):
# batch_voxels.append(voxels)
# batch_voxels =np.array(batch_voxels)# resolution * batchsize grid * grid * grid
# batch_labels = torch.from_numpy(np.array(batch_labels).astype(np.float32)) # batch size * 1
return batch_points, batch_labels
# append for inference KUDO
def getBatchsWithIdx(self, batch_size=1):
centers = []
for i in range(batch_size):
if self.currIdx >= self.num_batches:
break
centers.append([
self.xmin + self.spacing / 2 +
(self.currIdx // self.num_rows) * self.spacing,
self.ymin + self.spacing / 2 +
(self.currIdx % self.num_rows) * self.spacing
])
self.currIdx += 1
# print(centers)
if centers:
_, idx = self.tree.query(centers, k=self.k)
return self.points_and_features[
idx, :], self.labels[idx], np.array(idx)
else:
return None, None, None
def getBatchByIdx(self, batch_idx):
centers = [[
self.xmin + self.spacing / 2 +
(batch_idx // self.num_rows) * self.spacing, self.ymin +
self.spacing / 2 + (batch_idx % self.num_rows) * self.spacing
]]
_, idx = self.tree.query(centers, k=self.k)
return self.points_and_features[idx, :], self.labels[idx]
class Beacons:
def __init__(self, grid, beacon_grid=default_beacon_grid):
self.grid = grid
self.beacon_grid = beacon_grid
# self.beacons = [Beacon(location, count) for count, location in enumerate(self._uniform_deployment())]
# self.beacons = {count: Beacon(location, count) for count, location in enumerate(self._uniform_deployment())}
# self.beacons = {beac_tag: Beacon(location, beac_tag) for beac_tag, location in self._uniform_deployment()}
self.beacons = dict()
self.masks = []
self.map_closest_beacon = []
self.n = len(self.beacons)
# def _uniform_deployment(self):
# # return np.array(
# # list(itertools.product(np.linspace(0, self.grid.domain[0], self.beacon_grid[0] + 2)[1:-1],
# # np.linspace(0, self.grid.domain[1], self.beacon_grid[1] + 2)[1:-1])))
# # return np.array([default_nest_location])
# return zip([0,1], [default_nest_location, default_food_location])
# # return np.array(
# # list(itertools.product(np.linspace(default_nest_location[0]-2, default_nest_location[0]+2,
# # self.beacon_grid[0] + 2)[1:-1],
# # np.linspace(default_nest_location[1]-2, default_nest_location[1]+2,
# # self.beacon_grid[1] + 2)[1:-1]))
def update_amplitude(self):
for beac_tag in self.beacons:
self.beacons[beac_tag].amplitude()
def update_variance(self):
for beac_tag in self.beacons:
self.beacons[beac_tag].variance()
def update_beacon_configuration(self,
position_changed=True,
weights_changed=True,
just_initialized=False):
if position_changed:
self.update_masks()
if weights_changed:
self.update_neighbours_weights()
# TODO check why just_initialized condition is needed
if weights_changed and not just_initialized:
self.update_amplitude()
self.update_variance()
def move_step(self, W):
self.update_m_c_beacons(W)
self.update_locations()
self.update_beacon_configuration()
# def remove_step(self):
# self.remove_beacons()
# self.update_masks()
# self.update_neighbours_beacons()
# def switch_ant_beac_step(self,ants):
# ants.update_weights(self)
# # for ant in ants.ants:
# # if sum(ant.w) < 0.01:
# # self.beacons['aap'] = Beacon(ant.nt[1], 200)
# # ants.ants.remove(ant)
# test = 1
# def remove_beacons(self):
# weight_dict = self.check_weights(to_show = 'W',thres=threshold)
# ant_dict = self.check_ants()
# old_beacons = self.beacons.copy()
# for beac_tag in old_beacons:
# if beac_tag not in weight_dict and beac_tag not in ant_dict:
# del self.beacons[beac_tag]
def update_masks(self):
self.tree = KDTree(
[self.beacons[beac_tag].pt[1] for beac_tag in self.beacons])
self.map_closest_beacon = np.array([[
self.tree.query([self.grid.X[cy][cx], self.grid.Y[cy][cx]])[1]
for cx in range(len(self.grid.x))
] for cy in range(len(self.grid.y))])
self.masks = {
beac_tag: (self.map_closest_beacon == count) * 1
for count, beac_tag in enumerate(self.beacons)
}
for beac_tag in self.masks:
extended_mask = self.extend_mask(self.masks[beac_tag])
self.beacons[beac_tag].neigh = [
tag for tag in self.masks if True in (
extended_mask + self.masks[tag] >= 2) and tag != beac_tag
]
def update_neighbours_weights(self):
for beac_tag in self.masks:
self.beacons[beac_tag].neigh_weigh = [[
self.beacons[tag].w[0] for tag in self.beacons[beac_tag].neigh
], [
self.beacons[tag].w[1] for tag in self.beacons[beac_tag].neigh
]]
def evaporate_weights(self, rho=default_rho):
for beac_tag in self.beacons:
self.beacons[beac_tag].w *= (1 - rho)
def initialize_weights(self):
for beac_tag in self.beacons:
# replace n (nr beacons) by nr of foragers
# beacon.w = np.array([ampFactor * 1 / self.n, ampFactor * 1 / self.n])
# beacon.w = np.array([0., 0.])
self.beacons[beac_tag].w = np.array([0., 0.])
# self.beacons[beac_tag].w = np.array([offset, offset])
def update_m_c_beacons(self, W):
# for count, beacon in enumerate(self.beacons):
# beacon.mt[0] = beacon.mt[1]
# beacon.ct[0][0] = beacon.ct[0][1]
# beacon.ct[1][0] = beacon.ct[1][1]
# beacon.mt[1] = simps(simps(W * self.masks[count], self.grid.x), self.grid.y)
# beacon.ct[0][1] = simps(simps(W * self.grid.X * self.masks[count], self.grid.x), self.grid.y) / beacon.mt[1]
# beacon.ct[1][1] = simps(simps(W * self.grid.Y * self.masks[count], self.grid.x), self.grid.y) / beacon.mt[1]
for beac_tag in self.beacons:
self.beacons[beac_tag].mt[0] = self.beacons[beac_tag].mt[1]
self.beacons[beac_tag].ct[0][0] = self.beacons[beac_tag].ct[0][1]
self.beacons[beac_tag].ct[1][0] = self.beacons[beac_tag].ct[1][1]
self.beacons[beac_tag].mt[1] = simps(
simps(W * self.masks[beac_tag], self.grid.x), self.grid.y)
self.beacons[beac_tag].ct[0][1] = simps(
simps(W * self.grid.X * self.masks[beac_tag], self.grid.x),
self.grid.y) / self.beacons[beac_tag].mt[1]
self.beacons[beac_tag].ct[1][1] = simps(
simps(W * self.grid.Y * self.masks[beac_tag], self.grid.x),
self.grid.y) / self.beacons[beac_tag].mt[1]
def update_locations(self):
# for beacon in self.beacons:
# delta_ct = [(beacon.ct[0][1] - beacon.ct[0][0]), (beacon.ct[1][1] - beacon.ct[1][0])]
# delta_x = (delta_ct[0] / dt) - kappa * (beacon.pt[1][0] - beacon.ct[0][1]) + delta_ct[0] * \
# self.neigh_control_term(beacon)[0]
# delta_y = (delta_ct[1] / dt) - kappa * (beacon.pt[1][1] - beacon.ct[1][1]) + delta_ct[1] * \
# self.neigh_control_term(beacon)[1]
#
# beacon.pt[0] = beacon.pt[1]
# beacon.pt[1] = [beacon.pt[0][0] + bound(delta_x * dt),
# beacon.pt[0][1] + bound(delta_y * dt)]
for beac_tag in self.beacons:
if beac_tag not in [0, 1]:
delta_ct = [(self.beacons[beac_tag].ct[0][1] -
self.beacons[beac_tag].ct[0][0]),
(self.beacons[beac_tag].ct[1][1] -
self.beacons[beac_tag].ct[1][0])]
delta_x = (delta_ct[0] / dt) - kappa * (self.beacons[beac_tag].pt[1][0] - self.beacons[beac_tag].ct[0][1]) + \
delta_ct[0] * self.neigh_control_term(self.beacons[beac_tag])[0]
delta_y = (delta_ct[1] / dt) - kappa * (self.beacons[beac_tag].pt[1][1] - self.beacons[beac_tag].ct[1][1]) + \
delta_ct[1] * self.neigh_control_term(self.beacons[beac_tag])[1]
self.beacons[beac_tag].pt[0] = self.beacons[beac_tag].pt[1]
self.beacons[beac_tag].pt[1] = [
self.beacons[beac_tag].pt[0][0] + bound(delta_x * dt),
self.beacons[beac_tag].pt[0][1] + bound(delta_y * dt)
]
def neigh_control_term(self, beacon):
move_i = [0, 0]
for count_j in beacon.neigh:
j = self.beacons[count_j]
move_j = [j.pt[0][1] - j.pt[0][0], j.pt[1][1] - j.pt[1][0]]
if move_j[0] != 0:
move_i[0] += (1 / move_j[0]) * (
(j.ct[0][1] - j.ct[0][0]) / dt - kappa *
(j.pt[0][1] - j.ct[0][1]))
if move_j[1] != 0:
move_i[1] += (1 / move_j[1]) * (
(j.ct[1][1] - j.ct[1][0]) / dt - kappa *
(j.pt[1][1] - j.ct[1][1]))
return move_i
def fnc_ants_at_beacons(self, ants):
# for tag, beacon in enumerate(self.beacons):
# beacon.fnc_ants_at_beacon(ants,tag)
for beac_tag in self.beacons:
self.beacons[beac_tag].fnc_ants_at_beacon(ants)
@staticmethod
def extend_mask(mask):
return ((shift(mask, [1, 1], cval=0) + shift(mask, [1, -1], cval=0) +
shift(mask, [-1, 1], cval=0) + shift(mask, [-1, -1], cval=0) +
shift(mask, [0, -1], cval=0) + shift(mask, [0, 1], cval=0) +
shift(mask, [1, 0], cval=0) + shift(mask, [-1, 0], cval=0)) >=
1) * 1
def check_weights(self, to_check='W1', thres=0.):
if to_check == 'W1':
return {
beac_tag: self.beacons[beac_tag].w[0]
for beac_tag in self.beacons
if self.beacons[beac_tag].w[0] > thres
}
elif to_check == 'W2':
return {
beac_tag: self.beacons[beac_tag].w[1]
for beac_tag in self.beacons
if self.beacons[beac_tag].w[1] > thres
}
elif to_check == 'W':
return {
beac_tag:
self.beacons[beac_tag].w[0] + self.beacons[beac_tag].w[1]
for beac_tag in self.beacons
if self.beacons[beac_tag].w[0] > thres
or self.beacons[beac_tag].w[1] > thres
}
def check_ants(self, thres=0):
return {
beac_tag: self.beacons[beac_tag].ants_at_beacon
for beac_tag in self.beacons
if self.beacons[beac_tag].ants_at_beacon > thres
}
data = np.concatenate(
(np.random.randn(n // 2, m), np.random.randn(n - n // 2, m) + np.ones(m)))
queries = np.concatenate(
(np.random.randn(r // 2, m), np.random.randn(r - r // 2, m) + np.ones(m)))
print "dimension %d, %d points" % (m, n)
t = time.time()
T1 = KDTree(data)
print "KDTree constructed:\t%g" % (time.time() - t)
t = time.time()
T2 = cKDTree(data)
print "cKDTree constructed:\t%g" % (time.time() - t)
t = time.time()
w = T1.query(queries)
print "KDTree %d lookups:\t%g" % (r, time.time() - t)
del w
t = time.time()
w = T2.query(queries)
print "cKDTree %d lookups:\t%g" % (r, time.time() - t)
del w
T3 = cKDTree(data, leafsize=n)
t = time.time()
w = T3.query(queries)
print "flat cKDTree %d lookups:\t%g" % (r, time.time() - t)
del w
t = time.time()
pix = im.load()
pix_freqs = Counter(
[pix[x, y] for x in range(im.size[0]) for y in range(im.size[1])])
pix_freqs_sorted = sorted(pix_freqs.items(), key=lambda x: x[1])
pix_freqs_sorted.reverse()
# print(pix)
outImg = Image.new('RGB', im.size, color='white')
outFile = open(
filepath_bytes + filename + "_" + str(new_w) + "_" + str(new_h) +
file_text_ext, 'w')
i = 0
for y in range(im.size[1]):
for x in range(im.size[0]):
pixel = im.getpixel((x, y))
#print(pixel)
if (pixel[3] < 200):
outImg.putpixel((x, y), palette_rgb[0])
outFile.write("%x\n" % (0))
#print(i)
else:
index = pixel_tree.query(pixel[:3])[1]
outImg.putpixel((x, y), palette_rgb[index])
outFile.write("%x\n" % (index))
i += 1
outFile.close()
outImg.save(filepath_converted + filename + "_" + str(new_w) + "_" +
str(new_h) + file_picture_ext)
print 'Done: ' + filename + ', ' + str(iteration)
new_w = orig_w - width_inc
new_h = orig_h - height_inc
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.pose = None
self.waypoints = None
self.camera_image = None
self.waypoints_2d = None
self.waypoint_tree = None
self.lights = []
self.state2light = {0: 'Red', 1: 'Yellow', 2: 'Green', 4: 'Unknown'}
sub1 = rospy.Subscriber('/current_pose', PoseStamped, self.pose_cb)
sub2 = rospy.Subscriber('/base_waypoints', Lane, self.waypoints_cb)
'''
/vehicle/traffic_lights provides you with the location of the traffic light in 3D map space and
helps you acquire an accurate ground truth data source for the traffic light
classifier by sending the current color state of all traffic lights in the
simulator. When testing on the vehicle, the color state will not be available. You'll need to
rely on the position of the light and the camera image to predict it.
'''
sub3 = rospy.Subscriber('/vehicle/traffic_lights', TrafficLightArray,
self.traffic_cb)
sub6 = rospy.Subscriber('/image_color', Image, self.image_cb)
self.pose_lock = threading.RLock()
self.base_wp_lock = threading.RLock()
self.lights_lock = threading.RLock()
self.image_lock = threading.RLock()
self.bbox_lock = threading.RLock()
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
self.stop_line_positions = self.config['stop_line_positions']
rospy.loginfo("Length of self.stop_line_positions {0}".format(
len(self.stop_line_positions)))
self.stoplines_2d = [[stopline[0], stopline[1]]
for stopline in self.stop_line_positions]
self.stopline_tree = KDTree(self.stoplines_2d)
rospy.loginfo("Length of self.stoplines_2d {0}: {1}".format(
len(self.stoplines_2d), self.stoplines_2d))
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
self.bridge = CvBridge()
self.light_classifier = TLClassifier()
self.listener = tf.TransformListener()
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
self.prev_detected = TrafficLight.UNKNOWN
# darknet_ros message
sub5 = rospy.Subscriber('/darknet_ros/bounding_boxes', BoundingBoxes,
self.detected_bb_cb)
# Get simulator_mode parameter (1== ON, 0==OFF)
self.simulator_mode = rospy.get_param("/simulator_mode")
if int(self.simulator_mode) == 0:
self.cropped_tl_bb_pub = rospy.Publisher('/cropped_bb',
Image,
queue_size=1)
self.TL_BB_list = []
# atexit.register(self._ros_atexit)
# self.loop()
#
# def loop(self):
# rate = rospy.Rate(10) # 10hz
# while not rospy.is_shutdown():
# if None not in [self.pose, self.waypoints, self.camera_image, self.waypoints_2d, self.waypoint_tree]:
# self.update_waypoints()
# rate.sleep()
rospy.spin()
def detected_bb_cb(self, msg):
simulator_bb_size_threshold = 24 #px
site_bb_size_threshold = 40 #px
# min probability of detection
simulator_bb_probability = 0.75
site_bb_probability = 0.25
if int(self.simulator_mode) == 1:
prob_thresh = simulator_bb_probability
size_thresh = simulator_bb_size_threshold
else:
prob_thresh = site_bb_probability
size_thresh = site_bb_size_threshold
self.bbox_lock.acquire()
try:
self.TL_BB_list = []
for bb in msg.bounding_boxes:
if str(bb.Class
) == 'traffic light' and bb.probability >= prob_thresh:
if math.sqrt((bb.xmin - bb.xmax)**2 +
(bb.ymin - bb.ymax)**2) >= size_thresh:
self.TL_BB_list.append(bb)
if None not in [
self.pose, self.waypoints, self.camera_image,
self.waypoints_2d, self.waypoint_tree
] and len(self.TL_BB_list) > 0:
self.update_waypoints()
finally:
self.bbox_lock.release()
def _ros_atexit(self):
rospy.loginfo("Shutdown signal received")
def update_waypoints(self):
'''
Publish upcoming red lights at camera frequency.
Each predicted state has to occur `STATE_COUNT_THRESHOLD` number
of times till we start using it. Otherwise the previous stable state is
used.
'''
self.pose_lock.acquire()
try:
pose = self.pose
vehicle_idx = self.get_closest_waypoint(pose.pose.position.x,
pose.pose.position.y)
finally:
self.pose_lock.release()
self.base_wp_lock.acquire()
try:
waypoints = self.waypoints
finally:
self.base_wp_lock.release()
self.lights_lock.acquire()
try:
lights = self.lights
finally:
self.lights_lock.release()
self.image_lock.acquire()
try:
camera_image = self.camera_image
finally:
self.image_lock.release()
self.bbox_lock.acquire()
try:
TL_BB_list = self.TL_BB_list
finally:
self.bbox_lock.release()
light_wp, state = self.process_traffic_lights(pose, lights, TL_BB_list,
camera_image)
if self.state == TrafficLight.UNKNOWN or self.state != state:
self.state = state
self.state_count = 0
elif self.state == state:
self.state_count += 1
if self.state_count >= STATE_COUNT_THRESHOLD:
self.state_count = 0
self.last_state = self.state
self.state = state
behind = self.is_behind(waypoints, light_wp, vehicle_idx)
return_wp = light_wp if (
state == TrafficLight.RED
or state == TrafficLight.UNKNOWN) and not behind else -1
self.last_wp = return_wp
rospy.loginfo(
"{0} light at wp {1}, vehicle_wp: {2}, return_wp: {3} (behind: {4})"
.format(self.state2light[state], light_wp, vehicle_idx,
return_wp, behind))
self.upcoming_red_light_pub.publish(Int32(return_wp))
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
def pose_cb(self, msg):
self.pose_lock.acquire()
try:
self.pose = msg
finally:
self.pose_lock.release()
def waypoints_cb(self, lane):
self.base_wp_lock.acquire()
try:
self.waypoints = lane.waypoints
if self.waypoints_2d is None:
self.waypoints_2d = [[
waypoint.pose.pose.position.x,
waypoint.pose.pose.position.y
] for waypoint in self.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d)
finally:
self.base_wp_lock.release()
def traffic_cb(self, msg):
self.lights_lock.acquire()
try:
self.lights = msg.lights
finally:
self.lights_lock.release()
def image_cb(self, msg):
"""Identifies red lights in the incoming camera image and publishes the index
of the waypoint closest to the red light's stop line to /traffic_waypoint
Args:
msg (Image): image from car-mounted camera
"""
self.image_lock.acquire()
try:
self.camera_image = msg
self.has_image = True
finally:
self.image_lock.release()
def is_behind(self, waypoints, object_wp_idx, vehicle_wp_idx):
wps_len = len(waypoints)
if object_wp_idx == vehicle_wp_idx:
return False
if (object_wp_idx > vehicle_wp_idx):
if (vehicle_wp_idx + wps_len) - object_wp_idx < (object_wp_idx -
vehicle_wp_idx):
return True
else:
return False
else:
if (object_wp_idx + wps_len) - vehicle_wp_idx < (vehicle_wp_idx -
object_wp_idx):
return False
else:
return True
def get_closest_waypoint(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
wps_len = len(self.waypoints_2d)
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
closest_coord = self.waypoints_2d[closest_idx]
prev_coord = self.waypoints_2d[(wps_len + closest_idx - 1) % wps_len]
cl_vect = np.array(closest_coord)
prev_vect = np.array(prev_coord)
pose_vect = np.array([x, y])
val = np.dot(cl_vect - prev_vect, pose_vect - cl_vect)
if (val < 0.0):
closest_idx = (closest_idx + wps_len - 1) % wps_len
return closest_idx
def get_closest_stopline(self, x, y):
"""Identifies the closest path waypoint to the given position
https://en.wikipedia.org/wiki/Closest_pair_of_points_problem
Args:
pose (Pose): position to match a waypoint to
Returns:
int: index of the closest waypoint in self.waypoints
"""
closest_idx = self.stopline_tree.query([x, y], 1)[1]
#rospy.loginfo("get_closest_stopline x {0} y {1}".format(x, y))
#rospy.loginfo("Closest stop line {0}: x {1} y {2}".format(closest_idx, x, y))
#rospy.loginfo("Length of self.stoplines_2d {0}: {1}".format(len(self.stoplines_2d), self.stoplines_2d))
return closest_idx
def get_light_state(self, camera_image, light, tl_bb_list):
"""Determines the current color of the traffic light
Args:
light (TrafficLight): light to classify
Returns:
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
#if(not self.has_image):
# self.prev_light_loc = None
# return False
state = TrafficLight.UNKNOWN
if len(tl_bb_list) > 0:
cv_image = self.bridge.imgmsg_to_cv2(camera_image, "bgr8")
#if int(self.simulator_mode) == 0:
#Get classification
state = self.light_classifier.get_classification(
cv_image, tl_bb_list, self.simulator_mode)
return state
def process_traffic_lights(self, pose, lights, tl_bb_list, camera_image):
"""Finds closest visible traffic light, if one exists, and determines its
location and color
Returns:
int: index of waypoint closes to the upcoming stop line for a traffic light (-1 if none exists)
int: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
light = None
# List of positions that correspond to the line to stop in front of for a given intersection
# stop_line_positions = self.config['stop_line_positions']
light_idx = self.get_closest_stopline(pose.pose.position.x,
pose.pose.position.y)
# rospy.loginfo("Pose x {0} y {1}".format(pose.pose.position.x, pose.pose.position.y))
# rospy.loginfo("light x {0} y {1}".format(self.stop_line_positions[light_idx][0], self.stop_line_positions[light_idx][1]))
light_wp = self.get_closest_waypoint(
self.stop_line_positions[light_idx][0],
self.stop_line_positions[light_idx][1])
# rospy.loginfo("light wp {0}".format(light_wp))
light = self.lights[light_idx]
#TODO find the closest visible traffic light (if one exists)
if light is not None:
state = self.get_light_state(camera_image, light, tl_bb_list)
return light_wp, state
return -1, TrafficLight.UNKNOWN
def process(input_image, f, params, model_params, tf_sess, flist, kalmangroup):
oriImg = input_image # B,G,R order
if f % video_process == 0:
subset_all, all_peaks_all, t1, t2, t3 = predict(
oriImg, model_params, tf_sess)
else:
subset_all, all_peaks_all, t1, t2, t3 = predict(
oriImg,
model_params,
tf_sess,
lenimg=len(flist) if len(flist) != 0 else 1,
flist=flist)
canvas = input_image # B,G,R order
t4 = time.time()
flistnew = flist.copy()
tree = []
suspectpoint = []
if f != 0 and flist != []:
points = [[x[0], x[1]] for x in flist]
tree = KDTree(points)
for k in range(len(subset_all)):
subset = subset_all[k]
all_peaks = all_peaks_all[k]
newloc_all = []
if f % video_process == 0:
locx = 0
locy = 0
else:
locx = flist[k][0] - PAD
locy = flist[k][1] - PAD
for n in range(len(subset)):
loc = []
newloc = []
lost = 1
for i in [1, 0, 2]:
idx = int(subset[n][i])
if int(subset[n][i]) != -1:
location = list(
map(int,
all_peaks[i][int(idx - all_peaks[i][0][3])][0:2]))
location[0] = location[0] + locx
location[1] = location[1] + locy
newloc.append(location[0])
newloc.append(location[1])
location = tuple(location)
loc.append(location)
else:
lost = i
newloc.append(lost)
newloc_all.append(newloc)
for newloc in newloc_all:
if f != 0 and tree != []:
dis, index = tree.query([newloc[0], newloc[1]])
if dis > 20:
suspectpoint.append([newloc, dis, index])
else:
No = index
if detected[No] == 1:
flistnew[No] = Fuse(newloc, flistnew[No])
continue
else:
detected[No] = 1
flistnew[No] = newloc
else:
detected.append(1)
flistnew.append(newloc)
if suspectpoint != []:
for i in range(len(suspectpoint)):
if detected[suspectpoint[i][2]] == 1:
detected.append(1)
flistnew.append(suspectpoint[i][0])
else:
if suspectpoint[i][1] <= 80:
No = suspectpoint[i][2]
detected[No] = 1
flistnew[No] = suspectpoint[i][0]
if Kalman:
for i in range(len(flistnew)):
if i >= len(flist):
newKalman = CreateKalman()
newloc1 = [0, 0, 0, 0, 0, 0, 1]
n = 0
while (n <= 20):
newloc1 = Update(flistnew[i],
newKalman,
input_image.shape[1],
input_image.shape[0],
first=True)
n = n + 1
kalmangroup.append(newKalman)
flistnew[i] = newloc1
else:
flistnew[i] = Update(flistnew[i], kalmangroup[i],
input_image.shape[1],
input_image.shape[0])
for i in range(len(detected)):
item = detected[i]
if item <= -20:
flistnew.remove(flistnew[i])
detected.remove(detected[i])
if Kalman:
kalmangroup.remove(kalmangroup[i])
for i in range(len(detected)):
detected[i] = detected[i] - 1
t5 = time.time()
for i in range(len(flistnew)):
newloc = flistnew[i]
if detected[i] <= 0:
lost = newloc[-1]
if lost == 0:
loc = [(newloc[0], newloc[1]), (newloc[2], newloc[3])]
cv2.circle(canvas, loc[0], 4, colors[1], thickness=-1)
cv2.circle(canvas, loc[1], 4, colors[2], thickness=-1)
canvas = cv2.line(canvas, loc[0], loc[1], colors[1], 2)
elif lost == 2:
loc = [(newloc[0], newloc[1]), (newloc[2], newloc[3])]
cv2.circle(canvas, loc[0], 4, colors[0], thickness=-1)
cv2.circle(canvas, loc[1], 4, colors[1], thickness=-1)
canvas = cv2.line(canvas, loc[0], loc[1], colors[0], 2)
else:
loc = [(newloc[0], newloc[1]), (newloc[2], newloc[3]),
(newloc[4], newloc[5])]
cv2.circle(canvas, loc[0], 4, colors[1], thickness=-1)
cv2.circle(canvas, loc[1], 4, colors[0], thickness=-1)
cv2.circle(canvas, loc[2], 4, colors[2], thickness=-1)
canvas = cv2.line(canvas, loc[1], loc[0], colors[0], 2)
canvas = cv2.line(canvas, loc[0], loc[2], colors[1], 2)
cv2.putText(canvas, str(i), loc[0], font, 0.8, (255, 255, 255), 2)
flist = flistnew
t6 = time.time()
return canvas, t1, t2, t3, t4, t5, t6, flist, kalmangroup
class RtreeKDTree(object):
""" wrap scipy KDTree spatial index to look as much like Rtree class
as possible
"""
def __init__(self, stream, interleaved=False):
""" stream: an iterable, returning tuples of the form (id,[xmin,xmax,ymin,ymax],object)
for now, requires that xmin==xmax, and ymin==ymax
For now, only interleaved=False is supported.
"""
if interleaved:
raise Exception(
"No support for interleaved index. You must use xxyy ordering"
)
it = iter(stream)
self.points = []
self.data = []
for fid, xxyy, obj in it:
if xxyy[0] != xxyy[1] or xxyy[2] != xxyy[3]:
raise Exception(
"No support in kdtree for finite sized objects")
self.points.append([xxyy[0], xxyy[2]])
self.data.append((fid, obj))
self.refresh_tree()
def refresh_tree(self):
self.kdt = KDTree(self.points)
def nearest(self, rect, count):
xy = [rect[0], rect[2]]
dists, results = self.kdt.query(xy, k=count)
if count == 1:
dists = [dists]
results = [results]
return [self.data[r][0] for r in results]
def intersects(self, xxyy):
""" This should be made compatible with the regular RTree call...
"""
raise Exception("Intersects is not implemented in scipy.KDTree")
return []
def insert(self, feat_id, rect=None):
if rect[0] != rect[1] or rect[2] != rect[3]:
raise Exception("No support in kdtree for finite sized objects")
if rect is None:
raise Exception(
"Not sure what inserting an empty rectangle is supposed to do..."
)
self.points.append([rect[0], rect[2]])
self.data.append([feat_id, None])
self.refresh_tree()
def feat_id_to_index(self, feat_id):
for i in range(len(self.data)):
if self.data[i][0] == feat_id:
return i
raise Exception("feature id not found")
def delete(self, feat_id, rect):
index = self.feat_id_to_index(feat_id)
del self.data[index]
del self.points[index]
self.refresh_tree()
class SIR_Grid:
"""
A grid is a set of levels which in turn contain a set of nodes.
"""
def __init__(self, log_minimum, n, k):
#def __init__(self, n):
"""
"""
self.SIR_Levels = []
self.log_minimum = log_minimum
self.k = k
self.n = n
self.r0 = []
self.serial_interval = []
self.policies = []
self.best_policy_function = []
self.best_policy_indexes = []
def fill(self, r0, serial_interval, policies, fatality_rate,
healthcare_capacity, overcapacity_fatality_rate, time_to_vaccine):
shares = make_shares(self.n)
infected_levels = make_infected_levels(self.log_minimum, log(1 / r0),
self.k)
#shares = make_shares(self.n)
last_level = []
count_policies = [0]
self.r0 = r0
self.serial_interval = serial_interval
herd_immunity_thresholds = [[max(1 - 1 / r0, 0)], [0]]
available_policies = [SIR_Policy('none', r0, 0)]
data = []
self.best_policy_indexes = []
for policy in policies:
available_policies.append(
SIR_Policy(policy[0], policy[1], policy[2]))
herd_immunity_thresholds[0].append(max(1 - 1 / policy[1], 0))
herd_immunity_thresholds[1].append(policy[2] * time_to_vaccine)
count_policies.append(0)
self.policies = available_policies
for i, share in enumerate(reversed(shares)):
self.SIR_Levels.append(SIR_Levels(share))
current_level = self.SIR_Levels[i]
minimum_infected = infected_levels[0]
for j, infected in enumerate(infected_levels):
if infected <= share:
immune = share - infected
current_level.nodes.append(SIR_Node(immune, infected))
current_node = current_level.nodes[j]
current_node.fill(minimum_infected, last_level,
available_policies, serial_interval,
herd_immunity_thresholds, fatality_rate,
healthcare_capacity,
overcapacity_fatality_rate,
time_to_vaccine)
count_policies[
current_node.best_action_index] = count_policies[
current_node.best_action_index] + 1
data.append([immune, infected])
self.best_policy_indexes.append(
current_node.best_action_index)
if i + 1 < self.n:
current_level.make_cost()
last_level = current_level
self.best_policy_function = KDTree(data)
print(count_policies)
def find_best_policy(self, immune, infected):
"""
This function uses the KDTree object populated in the 'fill' method to approximate the cost minimizing
policy at the point (immune, infected). That policy will be based on the nearest node.
"""
distance, index = self.best_policy_function.query([immune, infected])
return (self.best_policy_indexes[index])
def plot(self, policies, colors):
"""
this function creates a plot of the optimal decision at each node in the state space using a color coded
arrow. The arrow indicates the start and end points of the epidemic under the action.
"""
ax = plt.axes(xlim=(0, 0.5), ylim=(0, 0.025))
ax.set_xlabel('share of population immune')
ax.set_ylabel('share of population infected')
ax.text(0.05, 0.0255, "none", color='g')
ax.text(0.15, 0.0255, "sustainable", color='b')
ax.text(0.25, 0.0255, "unsustainable", color='r')
for SIR_Level in self.SIR_Levels:
for node in SIR_Level.nodes:
action = node.best_action_index
start_arrow_x = node.actions[action].start_point.immune
start_arrow_y = node.actions[action].start_point.infected
length_arrow_x = node.actions[
action].end_point.immune - start_arrow_x
length_arrow_y = node.actions[
action].end_point.infected - start_arrow_y
ax.arrow(start_arrow_x,
start_arrow_y,
length_arrow_x,
length_arrow_y,
width=0.0005,
head_width=0.001,
color=colors[action])
plt.show()
def plot_with_projection(self, policies, colors, s0, st, it, date_range):
ax = plt.axes(xlim=(0, 0.5), ylim=(0, 0.025))
ax.set_xlabel('share of population immune')
ax.set_ylabel('share of population infected')
ax.text(0.05, 0.0255, "none", color='g')
ax.text(0.15, 0.0255, "sustainable", color='b')
ax.text(0.25, 0.0255, "unsustainable", color='r')
for SIR_Level in self.SIR_Levels:
for node in SIR_Level.nodes:
action = node.best_action_index
start_arrow_x = node.actions[action].start_point.immune
start_arrow_y = node.actions[action].start_point.infected
length_arrow_x = node.actions[
action].end_point.immune - start_arrow_x
length_arrow_y = node.actions[
action].end_point.infected - start_arrow_y
ax.arrow(start_arrow_x,
start_arrow_y,
length_arrow_x,
length_arrow_y,
width=0.0005,
head_width=0.001,
color=colors[action])
projection = self.project_epidemic(1, st / s0, it / s0, date_range)
ax.plot(projection['immune'], projection['infected'], lw=5, color='k')
plt.show()
def project_epidemic(self, s0, st, it, date_range):
"""
This function projects the path of the epidemic through the state space
"""
projected = pd.DataFrame(
columns=['date', 'policy', 'immune', 'infected', 'new infections'])
infected = it / s0
immune = 1 - st / s0
for date in date_range:
best_policy_index = self.find_best_policy(immune, infected)
policy_r0 = self.policies[best_policy_index].r0
policy = self.policies[best_policy_index].name
immune_after, infected_after = SIR_step(immune, infected,
policy_r0,
self.serial_interval)
new_infections = infected_after - infected + immune_after - immune
immune = immune_after
infected = infected_after
projected = projected.append(
{
'date': date,
'policy': policy,
'immune': immune * s0,
'infected': infected * s0,
'new infections': new_infections * s0
},
ignore_index=True)
return (projected)
class Mesh:
def __init__(self, filename, kdTree=False):
''' A mesh has already been generated on a file.'''
if filename is None:
self.vs = []
self.ws = []
self.cells = []
self.neighbours = {}
else:
self.filename = filename
f = open(filename)
lines = f.readlines()
if len(lines) == 0:
raise
self.dt = float(lines[1])
# preamble is two lines
self.n_char = len(lines) - 2
self.vs = []
self.ws = []
self.cells = []
# Only one inversion may occur in a meshfile
self.inversion = False
blocks = self.__split_blocks__(lines)
# it is now assumed that self.cells is a list of lists of quadrilaterals.
# this is not enforced yet
self.cells = []
# We create a place holder for strip 0. This cell has area 0; this can be used to test whether
# real stationary cells have been added or not
self.cells.append([Cell([0.0], [0.0])])
for block in blocks:
vs, ws = self.__build_arrays__(block)
self.vs.append(vs)
self.ws.append(ws)
self.__build_grid__(vs, ws)
if (kdTree == True): self.__build_tree__()
self.neighbours = {}
self.__build_neighbours__()
self.are_labels_drawn = True
def __split_blocks__(self, lines):
blocks = []
newblock = []
for line in lines[2:]:
if 'end' in line:
return blocks
if 'closed' in line or 'inversion' in line:
if len(newblock) % 2 != 0:
raise ValueError
blocks.append(list(newblock))
newblock = []
if 'inversion' in line:
if self.inversion == True:
raise ValueError
self.inversion = True
# demarcate where the inversion happens, so that later the number
# of strips before the inversion can be calculated
bound = len(blocks)
else:
newblock.append(line)
blocks.append(newblock)
# self bound is the upper boundary of the regular mesh. From this strip number onwards,
# strips are made from 'inversions'.
self.bound = 1
for i, block in enumerate(blocks):
if self.inversion == True and i == bound: break
nr_strips = len(block) / 2 - 1
self.bound += nr_strips
return blocks
def __build_arrays__(self, block):
'''Builds arrays of characteristics,
the v points and the w points of a characteristic
each have their own list.'''
vs = []
ws = []
for line in block[0::2]:
items = line.split()
data = [float(el) for el in items]
vs.append(data)
for line in block[1::2]:
items = line.split()
data = [float(el) for el in items]
ws.append(data)
return vs, ws
def __build_tree__(self):
''' build a KDTree out of the characteristic arrays. The grid points are stored
in self.data. Thhe same point may occur multiple times, as'''
points = []
tandem = zip(self.vs, self.ws)
for block in tandem:
blocks = zip(block[0], block[1])
for el in blocks:
coords = zip(el[0], el[1])
for coord in coords:
point = [coord[0], coord[1]]
points.append(point)
self.data = np.array(points)
self.tree = KDTree(self.data)
def __build_grid__(self, vs, ws):
''' Builds a list of lists of cells.'''
for i, v in enumerate(vs):
if i < len(vs) - 1:
m = min(len(vs[i]), len(vs[i + 1]))
cell = []
for j in range(m - 1):
cellv = [
vs[i][j], vs[i][j + 1], vs[i + 1][j + 1], vs[i + 1][j]
]
cellw = [
ws[i][j], ws[i][j + 1], ws[i + 1][j + 1], ws[i + 1][j]
]
try:
quad = Quadrilateral(cellv, cellw)
except ValueError:
print 'An error, probably a degeneracy in strip: ', i, ' cell: ', j
cell.append(quad)
self.cells.append(cell)
def __build_neighbours__(self):
'''Builds a dictionary of lists of cells, keyed by a tuple representing a point. The list contains all cells that
a meshpoint is part of. Only the quadrilaterals should be included in this neighbours list.'''
for i, cellist in enumerate(self.cells):
for j, cell in enumerate(cellist):
# this test ensures that only quadrilaterals are included in the neighbour list,
# because cell[0][0] is initialized with a Cell instances that has no area.
# After initialization cell[0][0] can have an area. i,j must run over the entire
# cell list as they are to become coordinates.
if cell.area > 0:
for p in range(4):
point = (cell.points[p][0], cell.points[p][1])
if point in self.neighbours:
self.neighbours[point].append([i, j])
else:
self.neighbours[point] = []
self.neighbours[point].append([i, j])
return
def __add_label__(self, cell, i, j, labelsize):
point = cell.centroid
t = ROOT.TText(point[0], point[1], str(i) + ',' + str(j))
t.SetTextAlign(22)
t.SetTextSize(labelsize)
return t
def insert_negative(self,
perimeter,
N=10000,
fiducial_distance=0.0,
exclusion_list=[]):
'''Replace the place holder for a stationary point by a Negative instance. If more than one stationary point is present,
extend the list.'''
cell = Negative(self.cells, perimeter, N, fiducial_distance,
exclusion_list)
if self.cells[0][0].area > 0:
self.cells[0].append(cell)
else:
self.cells[0][0] = cell
def insert_stationary(self, quadrilateral):
'''Quadrilateral needs to be a cell of type Quadrilateral.'''
if self.cells[0][0].area > 0:
self.cells[0].append(quadrilateral)
else:
self.cells[0][0] = quadrilateral
def deltat(self):
'''Time step used in building the grid.'''
return self.dt
def draw(self, plotlist, labelsize):
''' Draw all cells as list of TLine elements.'''
# ignore reversal bin
for i, celllist in enumerate(self.cells):
for j, cell in enumerate(celllist):
cell.draw(plotlist)
#i+1 because the reversal bin is not included, but enumerate makes i start at 0
#this only affects the text label
if self.are_labels_drawn == True:
t = self.__add_label__(cell, i, j, labelsize)
plotlist.append(t)
#make sure stationary bins comes on top
for cell in self.cells[0]:
cell.draw(plotlist)
def dimensions(self):
'''Returns minv, maxv, minw, maxw'''
vs = [
point[0] for celllist in self.cells[0:] for cell in celllist
for point in cell.points
]
ws = [
point[1] for celllist in self.cells[0:] for cell in celllist
for point in cell.points
]
return np.array([[min(vs), max(vs)], [min(ws), max(ws)]])
def bbox(self, i, j):
''' Give the bounding box of a given cell '''
return self.cells[i][j].bbox()
def query(self, point):
''' Give the result of a KDTree query'''
ret = self.tree.query(point)
return ret
def bin_from_point(self, point):
'''Give the bin the point belongs to. Return None if there is none.'''
# first check the stationary bins
for i, cell in enumerate(self.cells[0]):
if self.isinbin(0, i, point):
return [0, i]
# else walk the tree
close = self.tree.query(point, MAX_NEIGHBOURS)
nearests = self.data[close[1]]
for nearest in nearests:
ns = self.neighbours[(nearest[0], nearest[1])]
for n in ns:
if self.isinbin(n[0], n[1], point):
return n
return None
def isinbin(self, i, j, point):
''' True if point(v,w) is in bin i,j. False otherwise.'''
return self.cells[i][j].isPointInside(point)
def findvs(self, v):
''' Give a list of all bins that contain potential V.'''
l = []
for i, cells in enumerate(self.cells):
for j, cell in enumerate(cells):
box = cell.bbox()
if v >= box[0][0] and v < box[1][0]:
l.append([i, j])
return l
def checkSimple(self):
'''Checks whether Quadrilaterals are Simple'''
chksum = 0
for i, cells in enumerate(self.cells):
for j, cell in enumerate(cells):
chkCell = self.cells[i][j]
if chkCell.__class__.__name__ == 'Quadrilateral':
if chkCell.isSimple() == False:
chksum += 1
print i, j
return chksum
def checkSelfIntersection(self):
'''Checks for concave and self-intersecting Quadrilaterals'''
chkConcave = 0
chkSelfIntersect = 0
x_coords = []
y_coords = []
for i, cells in enumerate(self.cells):
for j, cell in enumerate(cells):
chkCell = self.cells[i][j]
if chkCell.__class__.__name__ == 'Quadrilateral':
if chkCell.isSimple() == False:
if chkCell.isSelfIntersecting() == False:
chkConcave += 1
else:
chkSelfIntersect += 1
print i, j
x_coords = np.append(x_coords, chkCell.centroid[0])
y_coords = np.append(y_coords, chkCell.centroid[1])
plt.scatter(x_coords, y_coords)
if chkSelfIntersect > 0:
plt.show()
return chkConcave, chkSelfIntersect
def checkSmallBins(self):
'''Checks for concave and self-intersecting Quadrilaterals'''
chkSum = 0
x_coords = []
y_coords = []
for i, cells in enumerate(self.cells):
for j, cell in enumerate(cells):
chkCell = self.cells[i][j]
if chkCell.__class__.__name__ == 'Quadrilateral':
if chkCell.isTooSmall() == True:
chkSum += 1
print i, j
x_coords = np.append(x_coords, chkCell.centroid[0])
y_coords = np.append(y_coords, chkCell.centroid[1])
plt.scatter(x_coords, y_coords)
if chkSum > 0:
plt.show()
return chkSum
####under construction
def mergeSmallBins(self, thresh=2e-3):
new_mesh = Mesh(None)
new_mesh.dt = self.dt
new_mesh.filename = self.filename
for i, cells in enumerate(self.cells):
if i == 0:
new_mesh.cells.append(self.cells[0])
else:
for j, cell in enumerate(cells):
chkCell = self.cells[i][j]
if j == 0:
new_mesh.cells.append([chkCell])
elif not chkCell.isTooSmall(threshold=thresh):
new_mesh.cells[i].append(chkCell)
elif chkCell.isTooSmall(threshold=thresh):
new_mesh.cells[i].append(
mergeQuads(chkCell, self.cells[i][-1]))
break
return new_mesh
def removeBadBins(self):
new_mesh = Mesh(None)
new_mesh.dt = self.dt # Added (MdK): 6/04/2017
for i, cells in enumerate(self.cells):
if i == 0:
new_mesh.cells.append(self.cells[0])
else:
for j, cell in enumerate(cells):
chkCell = self.cells[i][j]
if j == 0 and chkCell.isSelfIntersecting():
new_mesh.cells.append([Cell([0.], [0.])])
elif j == 0 and not chkCell.isSelfIntersecting():
new_mesh.cells.append([chkCell])
elif not chkCell.isSelfIntersecting():
new_mesh.cells[i].append(chkCell)
elif chkCell.isSelfIntersecting():
break
return new_mesh
def removeSmallBins(self, thresh=1e-5):
new_mesh = Mesh(None)
new_mesh.dt = self.dt
new_mesh.filename = self.filename
for i, cells in enumerate(self.cells):
if i == 0:
new_mesh.cells.append(self.cells[0])
else:
for j, cell in enumerate(cells):
chkCell = self.cells[i][j]
if j == 0 and chkCell.isTooSmall(threshold=thresh):
new_mesh.cells.append([Cell([0.], [0.])])
elif j == 0 and not chkCell.isTooSmall(threshold=thresh):
new_mesh.cells.append([chkCell])
elif not chkCell.isTooSmall(threshold=thresh):
new_mesh.cells[i].append(chkCell)
elif chkCell.isTooSmall(threshold=thresh):
break
return new_mesh
def removeWithRenewal(self, threshold=1e-8):
new_mesh = Mesh(None)
revname = self.filename.split('.')[0] + '.rev'
#new_mesh.filename = fn
f = open(revname, 'w')
f.write('<Mapping Type=\"Reversal\">\n')
flagged_cells = []
# end_strips = []
# limit_coords = []
for i, cells in enumerate(self.cells):
if i == 0 or i == 1:
new_mesh.cells.append(self.cells[i])
else:
for j, cell in enumerate(cells):
chkCell = self.cells[i][j]
if j == 0 and (chkCell.isSelfIntersecting()
or chkCell.isTooSmall(threshold)):
new_mesh.cells.append([Cell([0.], [0.])])
print i, j
elif j == 0 and not (chkCell.isSelfIntersecting()
or chkCell.isTooSmall(threshold)):
new_mesh.cells.append([chkCell])
elif not (chkCell.isSelfIntersecting()
or chkCell.isTooSmall(threshold)):
new_mesh.cells[i].append(chkCell)
if j == len(self.cells[i]) - 1:
#ind = np.argmin(distance.cdist([self.cells[1][k].centroid for k in range(len(self.cells[1]))],[chkCell.centroid],'euclidean'))
ind = np.argmin(
distance.cdist(
[C.centroid for C in self.cells[1]],
[chkCell.centroid], 'euclidean'))
f.write(repr(i) + ',0\t1,' + repr(ind) + '\t1.0\n')
# end_strips.append([i, j])
# limit_coords.append(ind)
elif (chkCell.isSelfIntersecting()
or chkCell.isTooSmall(threshold)):
if j != 0:
flagged_cells.append([i, j])
else:
print i, j
break
# f.write('</Mapping>')
# for i,xy in enumerate(end_strips):
# C = self.cells[xy[0]][xy[1]].centroid
# D = self.cells[1][limit_coords[i]].centroid
# plt.scatter(C[0],C[1])
# plt.scatter(D[0],D[1], color = 'r')
for coords in flagged_cells:
ind = np.argmin(
distance.cdist(
[x.centroid for y in (new_mesh.cells) for x in y],
[self.cells[coords[0]][coords[1]].centroid], 'euclidean'))
i, j = new_mesh.cellIndex(ind)
f.write(
repr(coords[0]) + ',0\t' + repr(i) + ',' + repr(j) + '\t1.0\n')
f.write('</Mapping>')
# plt.show()
return new_mesh
def cellIndex(self, index):
test_index = 0
for i, cells in enumerate(self.cells):
for j, cell in enumerate(cells):
if test_index == index:
return i, j
else:
test_index += 1
def ToXML(self, fn):
with open(fn, 'w') as f:
f.write('<Mesh>\n')
f.write('<TimeStep>')
f.write(str(self.dt))
f.write('</TimeStep>\n')
for i, cells in enumerate(self.cells):
f.write('<Strip>')
for j, cell in enumerate(cells):
for p in cell.points:
f.write("{:.12f}".format(p[0]))
f.write(' ')
f.write("{:.12f}".format(p[1]))
f.write(' ')
f.write('</Strip>\n')
f.write('</Mesh>\n')
def ToStat(self, fn):
with open(fn, 'w') as f:
f.write('<Stationary>\n')
for i, cells in enumerate(self.cells):
for j, cell in enumerate(cells):
if len(cell.points) == 4:
f.write('<Quadrilateral><vline>')
for p in cell.points:
f.write("{:.12f}".format(p[0]))
f.write(' ')
f.write('</vline><wline>')
for q in cell.points:
f.write("{:.12f}".format(q[1]))
f.write(' ')
f.write('</wline></Quadrilateral>\n')
f.write('</Stationary>')
def FromXML(self, fn, fromString=False):
'''Constructs a mesh from eithe a file, or a string if fromString == True.'''
if not fromString:
self.filename = fn
tree = ET.parse(fn)
root = tree.getroot()
else:
self.filename = ""
root = ET.fromstring(fn)
for ts in root.iter('TimeStep'):
self.dt = float(ts.text)
for str in root.iter('Strip'):
l = []
# An empty strip should be allowed for example as a place holder for stationary cells
if str.text is not None:
coords = [float(x) for x in str.text.split()]
if len(coords) % 8 != 0:
raise ValueError
n_chunck = len(coords) / 8
for i in range(0, n_chunck):
vs = []
ws = []
vs.append(coords[8 * i])
vs.append(coords[8 * i + 2])
vs.append(coords[8 * i + 4])
vs.append(coords[8 * i + 6])
ws.append(coords[8 * i + 1])
ws.append(coords[8 * i + 3])
ws.append(coords[8 * i + 5])
ws.append(coords[8 * i + 7])
quad = Quadrilateral(vs, ws)
l.append(quad)
self.cells.append(l)
self.__build_neighbours__()
