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)
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 features_by_centers(base, features):
try:
from scipy.spatial import KDTree
except ImportError:
from PathScripts.kdtree import KDTree
features = sorted(features,
key=lambda feature: getattr(base.Shape, feature).Surface.Radius,
reverse=True)
coordinates = [(cylinder.Surface.Center.x, cylinder.Surface.Center.y) for cylinder in
[getattr(base.Shape, feature) for feature in features]]
tree = KDTree(coordinates)
seen = {}
by_centers = {}
for n, feature in enumerate(features):
if n in seen:
continue
seen[n] = True
cylinder = getattr(base.Shape, feature)
xc, yc, _ = cylinder.Surface.Center
by_centers[xc, yc] = {cylinder.Surface.Radius: feature}
for coord in tree.query_ball_point((xc, yc), cylinder.Surface.Radius):
seen[coord] = True
cylinder = getattr(base.Shape, features[coord])
by_centers[xc, yc][cylinder.Surface.Radius] = features[coord]
return by_centers
def update(self, time_passed):
epsilon = 0.01
time_passed = time_passed if time_passed != 0 else epsilon
for i in range(self.steps_per_frame):
milliseconds = float(time_passed) / self.steps_per_frame
if _DEBUG_PRINT_MS:
self.mss += [milliseconds]
_mss = self.mss[10:]
if(len(_mss) % 100):
print "milliseconds=%s" % _mss[len(_mss) - 1]
if self.coord_changed:
point_matrix = numpy.array([d.coord for d in self.devices])
kdtree = KDTree(point_matrix)
for d in self.devices:
d._nbrs = []
for (i, j) in kdtree.query_pairs(self.radio_range):
self.devices[i]._nbrs += [self.devices[j]]
self.devices[j]._nbrs += [self.devices[i]]
for d in self.devices:
d._nbr_range = _Field()
d._nbr_vec = _Field()
for n in d._nbrs + [d]:
delta = n.coord - d.coord
d._nbr_range[n] = numpy.dot(delta, delta) ** 0.5
d._nbr_vec[n] = delta
self.coord_changed = False
for d in self.devices:
d.dostep(milliseconds)
def compute(self, dataset_pool):
parcels = self.get_dataset()
arr = self.get_dataset().sum_dataset_over_ids(dataset_pool.get_dataset('household'), constant=1)
coords = column_stack( (parcels.get_attribute("x_coord_sp"), parcels.get_attribute("y_coord_sp")) )
kd_tree = KDTree(coords, 100)
results = kd_tree.query_ball_tree(kd_tree, self.radius)
return array(map(lambda l: arr[l].sum(), results))
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]))
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))
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(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 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 findAffineTransform(test_srcs, ref_srcs, max_pix_tol = 2., min_matches_fraction = 0.8, invariantMap=None):
if len(test_srcs) < 3:
raise Exception("Test sources has less than the minimum value of points (3).")
if invariantMap is None:
invMap = InvariantTriangleMapping()
if len(ref_srcs) < 3:
raise Exception("Test sources has less than the minimum value of points (3).")
#generateInvariants should return a list of the invariant tuples for each asterism and
# a corresponding list of the indices that make up the asterism
ref_invariants, ref_asterisms = invMap.generateInvariants(ref_srcs, nearest_neighbors = 7)
ref_invariant_tree = KDTree(ref_invariants)
test_invariants, test_asterisms = invMap.generateInvariants(test_srcs, nearest_neighbors = 5)
test_invariant_tree = KDTree(test_invariants)
#0.03 is just an empirical number that returns about the same number of matches than inputs
matches_list = test_invariant_tree.query_ball_tree(ref_invariant_tree, 0.03)
matches = []
#t1 is an asterism in test, t2 in ref
for t1, t2_list in zip(test_asterisms, matches_list):
for t2 in np.array(ref_asterisms)[t2_list]:
matches.append(zip(t2, t1))
matches = np.array(matches)
invModel = invMap.matchTransform(ref_srcs, test_srcs)
nInvariants = len(matches)
max_iter = nInvariants
min_matches = min(10, int(nInvariants * min_matches_fraction))
bestM = ransac.ransac(matches, invModel, 1, max_iter, max_pix_tol, min_matches)
return bestM
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 compute(self, dataset_pool):
with logger.block(name="compute variable persons_within_DDD_of_parcel with DDD=%s" % self.radius, verbose=False):
results = None
with logger.block(name="trying to read cache file %s" % self.cache_file_name, verbose=False):
try:
results = self._load_results()
except IOError:
logger.log_warning("Cache file could not be loaded")
with logger.block(name="initialize datasets", verbose=False):
parcels = self.get_dataset()
arr = self.get_dataset().sum_dataset_over_ids(dataset_pool.get_dataset('household'), attribute_name="persons")
if not results:
with logger.block(name="initialize coords", verbose=False):
coords = column_stack( (parcels.get_attribute("x_coord_sp"), parcels.get_attribute("y_coord_sp")) )
with logger.block(name="build KDTree", verbose=False):
kd_tree = KDTree(coords, 100)
with logger.block(name="compute"):
results = kd_tree.query_ball_tree(kd_tree, self.radius)
with logger.block(name="cache"):
if not SimulationState().cache_directory_exists():
logger.log_warning("Cache does not exist and is created.")
SimulationState().create_cache_directory()
self._cache_results(results)
with logger.block(name="sum results", verbose=False):
return_values = array(map(lambda l: arr[l].sum(), results))
return return_values
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 compute(self, dataset_pool):
parcels = self.get_dataset()
arr = self.get_dataset().compute_variables(['parcel.aggregate(household.children>0)'], dataset_pool=dataset_pool)
coords = column_stack( (parcels.get_attribute("x_coord_sp"), parcels.get_attribute("y_coord_sp")) )
kd_tree = KDTree(coords, 100)
results = kd_tree.query_ball_tree(kd_tree, self.radius)
return array(map(lambda l: arr[l].sum(), results))
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)
def filter_obj_info( OBJ_INFO ):
#OBJ_info.append( [ X1[i], Y1[i],
# cRA10[i], cDEC10[i],
# MAG1[i],
# X2[j], Y2[j],
# cRA20[j], cDEC20[j],
# MAG2[j], SN1, SN2, NAME ] )
### filter on X1, X2, sn
orig_obj = OBJ_INFO[:]
OBJ_INFO = numpy.asarray( numpy.asarray( OBJ_INFO ).T[0:-1], dtype='float')
print OBJ_INFO[10]
print
print OBJ_INFO[11]
snv = OBJ_INFO[10]**2 + OBJ_INFO[11]**2
X = OBJ_INFO[0]
Y = OBJ_INFO[1]
tree = KDTree( numpy.asarray( [X,Y] ).T )
keep = []
for i in range(0, len(X)):
match_inds = tree.query_ball_point( [X[i],Y[i]], 3.0 )
if numpy.all( snv[i] >= snv[ match_inds ] ):
keep.append( orig_obj[i] )
print ' Original length: %d'%( len(orig_obj) )
print 'Collapsed length: %d'%( len(keep) )
return keep
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 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 _kill_duplicates(arr, minimum_distance=10):
"""
Attempts to eliminate garbage coordinates
arr is a 2D array of (npeaks)x4 coordinates.
0:1 is the x and y coordinates.
2 is the peak height
3 is the peak width
"""
import scipy
from scipy.spatial import KDTree
tree = KDTree(arr)
match_list = tree.query_ball_tree(tree,minimum_distance)
match_list=[list_item for list_item in match_list if len(list_item)>1]
chuck_list=[]
for match in match_list:
# compile the heights from the table
#heights=arr[match][:,2]
#best=np.argmax(heights)
match.remove(match[0])
chuck_list += match
keepers = range(arr.shape[0])
[keepers.remove(chuck) for chuck in chuck_list if chuck in keepers]
return arr[keepers]
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 get_picking_position_for_single_item_bin(self, bin_letter, cur_camera_R, cur_camera_t, limb, colorful=True, fit=False, threshold=0.01, crop=False, bin_bound=None, perturb_xform=None, return_bin_content_cloud=False, return_normal=False): ''' physical pre-condition: place the camera to the hard-coded tote-viewing position. return (x,y,z) coordinate in global frame of position that has something to suck. ''' bin_content_cloud = self.get_current_bin_content_cloud(bin_letter, cur_camera_R, cur_camera_t, limb, colorful=colorful, fit=fit, threshold=threshold, crop=crop, bin_bound=bin_bound, perturb_xform=perturb_xform) pos = self.find_max_density_3d(bin_content_cloud) bin_content_tree = KDTree(bin_content_cloud[:,0:3]) neighbor_idxs = bin_content_tree.query_ball_point(pos, r=0.002) print 'neighbor_idxs is of type', neighbor_idxs.__class__ if len(neighbor_idxs)==0: print 'no neighbors!' normal = [-1,0,0] else: pts = bin_content_tree[neighbor_idxs,0:3] normal = get_normal(pts) return_list = [pos] if return_bin_content_cloud: return_list += [bin_content_cloud] if return_normal: return_list += [normal] if len(return_list)==1: return return_list[0] else: return return_list
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)
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 extraCredit(tripLocations, startPolygon, endPolygon):
#indices is a list that should contain the indices of the trips in the tripLocations list
#which start in the startPolygon region and end in the endPolygon region
indices = []
#TODO: insert your code here. You should build the kdtree and use it to query the closest
# intersection for each trip
startpoint = []
endpoint = []
for item in tripLocations:
startpoint.append([item[0],item[1]])
endpoint.append([item[2],item[3]])
#print start
startTime = time.time()
rs = ((startRectangle[0][0] - startRectangle[0][1])**2 + (startRectangle[1][0] - startRectangle[1][1])**2)**0.5/float(2)
re = ((endRectangle[0][0] - endRectangle[0][1])**2 + (endRectangle[1][0] - endRectangle[1][1])**2)**0.5/float(2)
os = [(startRectangle[0][0] + startRectangle[0][1])/2,(startRectangle[1][0] + startRectangle[1][1])/2]
oe = [(endRectangle[0][0] + endRectangle[0][1])/2,(endRectangle[1][0] + endRectangle[1][1])/2]
Stree = KDTree(startpoint)
Etree = KDTree(endpoint)
startpoint = set(Stree.query_ball_point(os, rs))
endpoint = set(Etree.query_ball_point(oe, re))
s_e = list(startpoint.intersection(endpoint))
for point in s_e:
if point_inside_polygon(tripLocations[point][0],tripLocations[point][1],startPolygon) and point_inside_polygon(tripLocations[point][2],tripLocations[point][3],endPolygon):
indices.append(i)
return indices
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
def daves_super_saturate(atoms):
pos = atoms.get_positions()
tree = KDTree(atoms.get_positions())
list_tree = list(tree.query_pairs(1.430))
bondedTo = [ [] for i in xrange(len(atoms))]
for bond in list_tree:
bondedTo[bond[0]].append(bond[1])
bondedTo[bond[1]].append(bond[0])
Zs = atoms. get_atomic_numbers()
# figure out what needs a hydrogen atom
for iatom in xrange(len(atoms)):
nbonds = len( bondedTo[iatom] )
Z = Zs[iatom]
if (Z,nbonds) == (6,2):
print "we should add H to atom ", iatom
r0 = pos[iatom, :]
bond1 = pos[ bondedTo[iatom][0] , : ] - r0
bond2 = pos[ bondedTo[iatom][1], :] -r0
rH = -(bond1 + bond2)
rH = 1.09 * rH / np.linalg.norm(rH)
atoms.append(Atom('H', r0+rH ))
def compare_image(the_image):
"""Return the fraction of sources found in the original image"""
# pixel comparison is not good, doesn't work. Compare catalogs.
if isinstance(the_image, np.ma.MaskedArray):
full_algn = the_image.filled(fill_value=np.median(the_image))\
.astype('float32')
else:
full_algn = the_image.astype('float32')
full_algn[the_image == 0] = np.median(the_image)
import sep
bkg = sep.Background(full_algn)
thresh = 3.0 * bkg.globalrms
allobjs = sep.extract(full_algn - bkg.back(), thresh)
allxy = np.array([[obj['x'], obj['y']] for obj in allobjs])
from scipy.spatial import KDTree
ref_coordtree = KDTree(self.star_new_pos)
# Compare here srcs list with self.star_ref_pos
num_sources = 0
for asrc in allxy:
found_source = ref_coordtree.query_ball_point(asrc, 3)
if found_source:
num_sources += 1
fraction_found = float(num_sources) / float(len(allxy))
return fraction_found
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 jiggle_points(self):
max_y, min_y = self.y + self.max_height/2, self.y - self.max_height/2
more_to_do = (len(self.points) > 0)
while more_to_do:
for i in range(len(self.points)):
current = self.points.pop(0)
points_array = numpy.array([[p.x, p.y] for p in self.points])
kdtree = KDTree(points_array)
ids = kdtree.query_ball_point([current.x, current.y], self.min_dist)
# If there are any neighbours too near
counter = 0
while not ids == []:
rand = random.uniform(-self.min_dist, self.min_dist)
current.y = min(max_y, max(min_y, current.y + rand))
ids = kdtree.query_ball_point([current.x, current.y], self.min_dist)
counter += 1
if counter > 100: break
self.points.append(current)
for i in range(len(self.points)):
current = self.points.pop(0)
points_array = numpy.array([[p.x, p.y] for p in self.points])
kdtree = KDTree(points_array)
ids = kdtree.query_ball_point([current.x, current.y], self.min_dist)
self.points.append(current)
more_to_do = (not ids == [])
if more_to_do: break
class WaypointUpdater(object):
def __init__(self):
rospy.init_node('waypoint_updater')
rospy.Subscriber('/current_pose',
PoseStamped,
self.pose_cb,
queue_size=2)
rospy.Subscriber('/base_waypoints',
Lane,
self.waypoints_cb,
queue_size=8)
# TODO: Add a subscriber for /traffic_waypoint and /obstacle_waypoint below
rospy.Subscriber('/traffic_waypoint', Int32, self.traffic_cb)
#rospy.Subscriber('/obstacle_waypoint', Int32, self.obstacle_cb, queue_size=1)
self.final_waypoints_pub = rospy.Publisher('final_waypoints',
Lane,
queue_size=1)
# TODO: Add other member variables you need below
self.base_lane = None
self.pose = None
self.waypoints_2d = None
self.waypoint_tree = None
self.stopline_wp_idx = -1
#rospy.spin()
self.decelerate_count = 0
self.loop()
# Infinite loop
def loop(self):
rate = rospy.Rate(PUB_FREQUENCY)
while not rospy.is_shutdown():
if self.pose and self.base_lane and self.waypoint_tree:
self.publish_waypoints()
rate.sleep()
def pose_cb(self, msg):
# TODO: Implement
self.pose = msg
#rospy.logwarn("Pose")
def waypoints_cb(self, waypoints):
# TODO: Implement
self.base_lane = waypoints
#rospy.logwarn("Waypoints in Baselane")
if not self.waypoints_2d:
self.waypoints_2d = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in waypoints.waypoints]
#rospy.logwarn("Waypoints in waypoint treed{}".format(self.waypoints_2d))
self.waypoint_tree = KDTree(self.waypoints_2d)
#rospy.logwarn("Waypoints in waypoint tree")
def traffic_cb(self, msg):
# TODO: Callback for /traffic_waypoint message. Implement
#if self.stopline_wp_idx != msg.data:
rospy.logwarn(
"LIGHT: new stopline_wp_idx={}, old stopline_wp_idx={}".format(
msg.data, self.stopline_wp_idx))
self.stopline_wp_idx = msg.data
def obstacle_cb(self, msg):
# TODO: Callback for /obstacle_waypoint message. We will implement it later
# I am not using it in this project
pass
# Get the closest waypoint index
def get_closest_waypoint_idx(self):
# Current pose (x,y) and index
x = self.pose.pose.position.x
y = self.pose.pose.position.y
closest_idx = self.waypoint_tree.query([x, y], 1)[1]
# Check if point is ahead or behind the car
closest_coord = self.waypoints_2d[closest_idx]
prev_coord = self.waypoints_2d[closest_idx - 1]
# Equation for hyperplane through closest_coords
cl_vect = np.array(closest_coord)
prev_vect = np.array(prev_coord)
pos_vect = np.array([x, y])
val = np.dot(cl_vect - prev_vect, pos_vect - cl_vect)
if val > 0:
closest_idx = (closest_idx + 1) % len(self.waypoints_2d)
return closest_idx
# Publishes the final lane
def publish_waypoints(self):
final_lane = self.generate_lane()
self.final_waypoints_pub.publish(final_lane)
def generate_lane(self):
lane = Lane()
# Get the closest and farthest index
closest_idx = self.get_closest_waypoint_idx()
farthest_idx = closest_idx + LOOKAHEAD_WPS
base_waypoints = self.base_lane.waypoints[closest_idx:farthest_idx]
rospy.logwarn("closest index :{} and stopline index:{}".format(
closest_idx, self.stopline_wp_idx))
# If no traffic light was detected, publish the base_waypoints as it is
if (self.stopline_wp_idx
== -1) or (self.stopline_wp_idx >= farthest_idx):
lane.waypoints = base_waypoints
rospy.logwarn("No Change")
else:
rospy.logwarn("Reduce speed")
lane.waypoints = self.decelerate_waypoints(base_waypoints,
closest_idx)
return lane
def decelerate_waypoints(self, waypoints, closest_idx):
temp = []
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.:
vel = 0
p.twist.twist.linear.x = min(vel, wp.twist.twist.linear.x)
temp.append(p)
self.decelerate_count += 1
if (self.decelerate_count % 5) == 0:
size = len(waypoints) - 1
vel_start = temp[0].twist.twist.linear.x
vel_end = temp[size].twist.twist.linear.x
rospy.logwarn("DECEL: vel[0]={:.2f}, vel[{}]={:.2f}".format(
vel_start, size, vel_end))
return temp
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 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
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector_site')
self.pose = None
self.waypoints = None
self.camera_image = None
self.lights = []
self.xycoords_orig_waypoints = None
self.kdtree_orig_waypoints = None
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
self.bridge = CvBridge()
self.light_classifier_site = TLClassifier()
#classify a blank black image at first, since the first image takes too long to classify
self.light_classifier_site.get_classification(np.zeros((600, 800, 3)))
self.listener = tf.TransformListener()
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
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.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
rospy.spin()
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, waypoints):
self.waypoints = waypoints
if not self.xycoords_orig_waypoints:
self.xycoords_orig_waypoints = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in waypoints.waypoints]
self.kdtree_orig_waypoints = KDTree(self.xycoords_orig_waypoints)
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
state = TrafficLight.UNKNOWN
if self.pose and self.kdtree_orig_waypoints:
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
or state == TrafficLight.YELLOW) else -1
self.last_wp = light_wp
self.upcoming_red_light_pub.publish(Int32(light_wp))
#rospy.loginfo("Red light at " + str(light_wp))
else:
self.upcoming_red_light_pub.publish(Int32(self.last_wp))
#rospy.loginfo("published previous signal")
self.state_count += 1
def get_closest_waypoint(self, xcoord, ycoord):
"""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
"""
id_closest_waypoint = self.kdtree_orig_waypoints.query(
[xcoord, ycoord], 1)[1]
return id_closest_waypoint
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")
state = self.light_classifier_site.get_classification(cv_image)
return 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_traffic_light = None
if (self.pose):
xpos = self.pose.pose.position.x
ypos = self.pose.pose.position.y
id_wp_car_position = self.get_closest_waypoint(xpos, ypos)
# 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']
#min_delta_wpid_car_line = len(self.waypoints.waypoints)
min_delta_wpid_car_line = 100 #using only 100 waypoints limits load on camera/classification
for i, light in enumerate(self.lights):
stop_line_xycoord = stop_line_positions[i]
id_wp_stop_line_i = self.get_closest_waypoint(
stop_line_xycoord[0], stop_line_xycoord[1])
delta_wpid_car_line = id_wp_stop_line_i - id_wp_car_position
if delta_wpid_car_line >= 0 and delta_wpid_car_line < min_delta_wpid_car_line:
id_wp_closest_stop_line = id_wp_stop_line_i
min_delta_wpid_car_line = delta_wpid_car_line
closest_traffic_light = light
if closest_traffic_light:
state = self.get_light_state(closest_traffic_light)
return id_wp_closest_stop_line, 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.waypoint_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.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.counter = 0
#self.light_list = []
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
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
"""
# 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)
"""
# Save Images
""" cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
resized_image = cv2.resize(cv_image, (320, 240))
cv2.imwrite("images/img-{0}.png".format(self.counter),resized_image)
img = []
img.append("img-{0}.png".format(self.counter))
img.append(light.state)
#self.light_list.append(img)
self.counter += 1
with open("training/output.csv", "a") as f:
writer = csv.writer(f)
writer.writerow(img)
f.close() """
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_wp_idx = self.get_closest_waypoint(self.pose.pose.position.x, self.pose.pose.position.y)
# 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
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 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
def calculate_surface_points(receptor,
ligand,
num_points,
rec_translation,
surface_density,
seed=STARTING_POINTS_SEED,
has_membrane=False,
num_sphere_points=100):
"""Calculates the position of num_points on the surface of the given protein"""
if num_points < 0:
raise SetupError(
"Invalid number of points to generate over the surface")
receptor_atom_coordinates = receptor.representative(has_membrane)
distances_matrix_rec = distance.pdist(receptor_atom_coordinates)
receptor_max_diameter = np.max(distances_matrix_rec)
distances_matrix_lig = distance.pdist(ligand.representative())
ligand_max_diameter = np.max(distances_matrix_lig)
surface_distance = ligand_max_diameter / 4.0
# Surface
pdb_file_name = Path(
receptor.structure_file_names[receptor.representative_id])
molecule = parsePDB(pdb_file_name).select('protein or nucleic')
if has_membrane:
pdb_no_membrane = str(
pdb_file_name.absolute().parent /
f"{pdb_file_name.stem}_no_membrane{pdb_file_name.suffix}")
writePDB(pdb_no_membrane, molecule)
surface = molecule.select('protein and surface or nucleic and name P')
coords = surface.getCoords()
# SASA
if num_points == 0:
if has_membrane:
structure = freesasa.Structure(pdb_no_membrane)
else:
structure = freesasa.Structure(str(pdb_file_name))
result = freesasa.calc(structure)
total_sasa = result.totalArea()
num_points = ceil(total_sasa / surface_density)
# Surface clusters
if len(coords) > num_points:
# Extremely important to set seed in order to get reproducible results
np.random.seed(seed)
surface_clusters = kmeans2(data=coords,
k=num_points,
minit='points',
iter=100)
surface_centroids = surface_clusters[0]
else:
surface_centroids = coords
# Create points over the surface of each surface cluster
sampling = []
for sc in surface_centroids:
sphere_points = np.array(points_on_sphere(num_sphere_points))
surface_points = sphere_points * surface_distance + sc
sampling.append(surface_points)
# Filter out not compatible points
centroids_kd_tree = KDTree(surface_centroids)
for i_centroid in range(len(sampling)):
# print('.', end="", flush=True)
centroid = surface_centroids[i_centroid]
# Search for this centroid neighbors
centroid_neighbors = centroids_kd_tree.query_ball_point(centroid,
r=20.)
# For each neighbor, remove points too close
for n in centroid_neighbors:
points_to_remove = []
if n != i_centroid:
for i_p, p in enumerate(sampling[i_centroid]):
if np.linalg.norm(
p - surface_centroids[n]) <= surface_distance:
points_to_remove.append(i_p)
points_to_remove = list(set(points_to_remove))
sampling[i_centroid] = [sampling[i_centroid][i_p] \
for i_p in range(len(sampling[i_centroid])) if i_p not in points_to_remove]
s = []
for points in sampling:
s.extend(points)
# Final cluster of points
if len(s) > num_points:
# Extremely important to set seed in order to get reproducible results
np.random.seed(seed)
s_clusters = kmeans2(data=s, k=num_points, minit='points', iter=100)
s = s_clusters[0]
for p in s:
p += rec_translation
return s, receptor_max_diameter, ligand_max_diameter
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_light_pub = rospy.Publisher('/traffic_waypoint', Int32, queue_size=1)
self.upcoming_red_light_pub = rospy.Publisher('/red_traffic_waypoint', Int32, queue_size=1)
self.classifier_state_pub = rospy.Publisher('/class_state', 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.cl_stat=0
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
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.
'''
self.upcoming_light_pub.publish(Int32(light_wp))
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.UNKNOWN:
self.cl_stat =9
else:
self.cl_stat =1
self.classifier_state_pub.publish(Int32(self.cl_stat))
light_wp = light_wp if state == TrafficLight.RED else -1
#rospy.logwarn('red light waypoint: {}'.format(light_wp))
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
#x = self.pose.pose.position.x
#y = self.pose.pose.position.y
closest_idx = self.waypoint_tree.query([x,y],1)[1]
#rospy.logwarn(closest_idx)
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)
"""
## ground truth
# 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
result = self.light_classifier.get_classification(cv_image)
#rospy.logwarn('ground truth/prediction: {0}/{1}'.format(COLOR_NAME_MAPPING[light.state], COLOR_NAME_MAPPING[result]))
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)
"""
#light = None
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 = self.get_closest_waypoint(self.pose.pose)
car_wp_idx = self.get_closest_waypoint(self.pose.pose.position.x, self.pose.pose.position.y)
diff = len(self.waypoints.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_wp_idx
#rospy.logwarn('distance between ego and closest light is :{}'.format(d))
if d>=0 and d<diff:
diff=d
closest_light=light
line_wp_idx=temp_wp_idx
#rospy.logwarn(closest_light.state)
#TODO find the closest visible traffic light (if one exists)
if closest_light:
state = self.get_light_state(closest_light)
#rospy.logwarn('closest light state: {}'.format(state))
#rospy.logwarn('closest stop line index: {}'.format(line_wp_idx))
#rospy.logwarn('closest ego index: {}'.format(car_wp_idx))
return line_wp_idx, state
#self.waypoints = None
return -1, TrafficLight.UNKNOWN
def cal_all_mapping_action(writer, sheetName, e_df, l_df):
"""
Main function to calcualte mapping action:
First using kdtree to get the nearest 20 points to search,
and calculate the 20 points for the possible loss value.
Finally get the minmum loss from the possible candidate above.
"""
# Parse the whole dataFrame into two list of 3Darrays
print('INTO dataE')
with ThreadPoolExecutor(max_workers=MAX_THREAD_NUM) as exe:
dataE = list(
exe.map(parse_row_to_list_of_3Darray, e_df.values, chunksize=20))
print('Finish dataE')
dataL = l_df.values
# Parse the last item, end position, to find the nearest 20 actions,
# to calculate possible loss values.
print('INTO dataEndPos')
parse_end_posE = functools.partial(parse_end_pos_to_3Darray, idx=E_END_POS)
parse_end_posL = functools.partial(parse_end_pos_to_3Darray, idx=L_END_POS)
with ThreadPoolExecutor(max_workers=MAX_THREAD_NUM) as exe:
dataEndPosE = list(exe.map(parse_end_posE, e_df.values, chunksize=20))
with ThreadPoolExecutor(max_workers=MAX_THREAD_NUM) as exe:
dataEndPosL = list(exe.map(parse_end_posL, l_df.values, chunksize=200))
print('Finish dataEndPos')
print('INTO rowCandidate')
# Using KDTree to find nearnest 20 end positions.
kdL = KDTree(dataEndPosL)
print('Finish rowCandidate')
# rowCandidate stores 20 indices of expert dataFrame, which are 20 nearest
# end positions.
rowCandidate = []
for rowData in dataEndPosE:
_, idx = kdL.query(x=rowData, k=50)
rowCandidate.append(idx)
# Using ThreadPoolExecutor will decreasing the data comparing time while looking up table, dataL.
# And without the ordering problem.
print('INTO mapData')
partial_cal_action = functools.partial(cal_mapping_action, dataL=dataL)
with ThreadPoolExecutor(max_workers=MAX_THREAD_NUM) as exe:
mapData = list(exe.map(partial_cal_action, dataE, rowCandidate))
# mapData = cal_angLoss(mapData)
data = np.array(mapData)
# data = smooth_mapping_data(data, dataE, dataL, rowCandidate)
# data = np.array(mapData)
print('Finish mapData')
# After the construction of all action mapping and loss calculation,
# I use pandas to write into excel.
print('INTO Excel')
header = []
for i, col in enumerate(list(l_df.columns)):
if (i % L_NUM_IN_TUPLE == 0):
# e.g [joint 0 pos] [joint 0 posSensor] [joint 1 pos] [joint 1 posSensor]
# so need to pick i % L_NUM_IN_TUPLE == 0.
header.append(col)
header[0] = 'learner num'
header.append('loss')
header.append('ang loss')
df = pd.DataFrame(data=data, columns=header)
df.to_excel(writer, sheet_name=sheetName, index=True)
# set layout
(_, max_col) = df.shape
worksheet = writer.sheets[sheetName]
worksheet.set_column(0, max_col, 15)
print('Finish Excel')
ent_indx += 1
entities.append(ent_indx)
pose_map[ent_indx] = []
pose_map[ent_indx].append([tuple(f) for f in frame[x]])
print('\t> Append Entity ID: ', ent_indx)
# Construct 2D-KD Tree Mapping of Joints
X = []
Y = []
lookup = dict()
for y in range(18):
X.append(frame[x][y][0])
Y.append(frame[x][y][1])
lookup[tuple(frame[x][y])] = ent_indx
kd_frame.append(KDTree(zip(X, Y)))
kp_lookup.append(lookup)
# pp = pprint.PrettyPrinter()
# pp.pprint(kp_lookup)
elif len(frame) == len(entities): # Same count as previous frame
print('No changes in entity counts...')
# Map Candidate from Previous Frame and Build Lookup Dictionary
for x in range(len(frame)):
# Query Candidate Joint
cand = []
for y in range(18):
query = kd_frame[x].query(frame[x][y], len(frame))
cand.append(zip(query[0], query[1]))
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]
#Creating a K-Dimension Binary search tree
self.waypoints_tree = KDTree(self.waypoints_2D)
def fit(self, X, grid=1.):
"""
fit the Means-Shift Model
:param X:
:param grid:
:return:
"""
centers = self.genMeshGridPoints(X, grid)
nc = len(centers)
kdtree = KDTree(X, leafsize=8)
centers_update = copy.deepcopy(centers)
for i in range(self.max_iter):
print('iter: ', i)
delete_idx = []
rnn_nums = [0] * nc
for j in range(nc):
query = centers[j, :]
neighbot_index = kdtree.query_ball_point(query, r=self.r)
rnn_nums[j] = len(neighbot_index)
if len(neighbot_index) > self.lessK:
# if the current center has neighbors, then recompute the center
center = np.mean(X[neighbot_index, :], axis=0)
centers[j, :] = center
else:
# if the current center has no neighbors, record the index and prepare to delete it
delete_idx.append(j)
# delete some centers who has no neighbors
if delete_idx:
centers = np.delete(centers, delete_idx, axis=0)
centers_update = np.delete(centers_update, delete_idx, axis=0)
rnn_nums = np.delete(np.array(rnn_nums), delete_idx, axis=0)
nc = len(centers)
# NMS, Non-Maximum Suppression
# 如果两个中心点坐标之差很小，则剔除radiusNN点较少的中心点
delete_idx = set()
for jj in range(nc):
diff_c = centers - centers[jj, :]
diff = np.sum(diff_c**2, axis=1)
diff[jj] = 1e3 # ensure the current center will not be delete
nms_idx = np.array(range(nc))[
diff < self.
nms] # get the point indexs that around the current center
for k in nms_idx:
if rnn_nums[k] < rnn_nums[
jj]: # record and delete the index whose neighbor'size is less
delete_idx.add(k)
if delete_idx:
delete_idx = list(delete_idx)
centers = np.delete(centers, delete_idx, axis=0)
centers_update = np.delete(centers_update, delete_idx, axis=0)
nc = len(centers)
show = 1
if show:
plt.figure()
plt.scatter(datas[:, 0], datas[:, 1], s=5, color='r')
plt.scatter(centers[:, 0], centers[:, 1], c='blue', s=8)
ax = plt.gca()
plot_args = {
'fc': 'None',
'lw': 2,
'edgecolor': 'b',
'ls': ':'
}
for j in range(nc):
circle = Circle(centers[j, :], radius=self.r, **plot_args)
ax.add_patch(circle)
plt.show()
# determine whether the iteration terminates or not
diff_sum = np.sum(np.fabs(centers_update - centers))
if diff_sum < self.tolerance:
print('centers do not change, break at iter:{}!'.format(i))
break
else:
centers_update = copy.deepcopy(centers)
# TODO: need NMS again
self.n_clusters, self.centers = len(centers), copy.deepcopy(centers)
def kdtree(self):
if 'kdtree' in self._cache:
return self._cache.get('kdtree')
with self._cache:
kdtree = KDTree(self.vertices.view(np.ndarray))
return self._cache.set('kdtree', kdtree)
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.has_image = False
self.pose = None
self.waypoints = None
self.waypoints_2d = None
self.waypoints_tree = None
self.camera_image = None
self.lights = []
rospy.Subscriber('/current_pose', PoseStamped, self.cb_pose)
rospy.Subscriber('/base_waypoints', Lane, self.cb_waypoints)
'''
/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.cb_traffic)
rospy.Subscriber('/image_color', Image, self.cb_image)
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
self.state_mapper = {
TrafficLight.UNKNOWN: 'UNKNOWN',
TrafficLight.RED: 'RED',
TrafficLight.YELLOW: 'YELLOW',
TrafficLight.GREEN: 'GREEN',
}
rospy.spin()
def cb_pose(self, msg):
self.pose = msg
def cb_waypoints(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 cb_traffic(self, msg):
self.lights = msg.lights
def cb_image(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('Caught traffic light WP at index {} with state {}'.format(
light_wp, self.state_mapper[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
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:
x : position to match a waypoint to in the x direction
y : position to match a waypoint to in the y direction
Returns:
int: index of the closest waypoint in self.waypoints
"""
return self.waypoints_tree.query([x, y], 1)[1]
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
return light.state
# cv_image = None
# 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)
"""
if not self.waypoints_tree:
# cant do any logic about traffic lights positions
return -1, TrafficLight.UNKNOWN
closest_light = None
wp_index_line = 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:
wp_index_car = self.get_closest_waypoint(
self.pose.pose.position.x, self.pose.pose.position.y)
lookahead = len(self.waypoints.waypoints)
for i, (light, line) in enumerate(zip(self.lights, stop_line_positions)):
# since there are just a few traffic lights wps doesnt make any
# sense to use a kdtree here, just iterate and find a closest wp
# in the the neighborhood. also, the problem of getting a wp
# behind the car here is less important as there are much less
# wps and more spaced apart
wp_index_tmp = self.get_closest_waypoint(line[0], line[1])
diff = wp_index_tmp - wp_index_car
if 0 <= diff < lookahead:
lookahead = diff
closest_light = light
wp_index_line = wp_index_tmp
if closest_light:
state = self.get_light_state(closest_light)
return wp_index_line, state
return -1, TrafficLight.UNKNOWN
def waypoints_cb(self, waypoints):
if not self.base_waypoints:
self.base_waypoints = waypoints
self.base_waypoint_count = len(waypoints.waypoints)
waypoints_xy = [[waypoint.pose.pose.position.x, waypoint.pose.pose.position.y] for waypoint in waypoints.waypoints]
self.waypoints_tree = KDTree(waypoints_xy)
def test_kdtree_comparisons():
# Regression test: node comparisons were done wrong in 0.12 w/Py3.
nodes = [KDTree.node() for _ in range(3)]
assert_equal(sorted(nodes), sorted(nodes[::-1]))
def build_correspondence(VS, FS, VT, FT, maxind, thres, FileNameTo_Save):
"""
build correspondence using the proximity and face normals of source and target meshes
:param VS: deformed source mesh matched with target nS * 3
:param FS: Triangle indices of source mesh mS * 3
:param VT: Target mesh nT * 3
:param FT: Triangle indices of target mesh mT * 3
:param maxind: Maximum correspondence
:param thres: Distance threshold for correspondence
:param FileNameTo_Save: string for speed up code
:return: corres mT * # of correspondence for each triangles of target mesh
"""
if os.path.exists(FileNameTo_Save):
corres = load_pickle_file(FileNameTo_Save)
return corres
else:
TS, NS, VS4, FS4 = v4_normal(VS, FS)
TT, NT, VT4, FT4 = v4_normal(VT, FT)
VS_C = np.zeros((FS.shape[0], 3))
VT_C = np.zeros((FT.shape[0], 3))
for i in range(0, FT.shape[0]):
VT_C[i, :] = np.mean(VT[FT[i, :] - 1, :], axis=0)
for i in range(0, FS.shape[0]):
VS_C[i, :] = np.mean(VS[FS[i, :] - 1, :], axis=0)
S_tree = KDTree(VS_C)
T_tree = KDTree(VT_C)
corres1 = -np.ones((FT.shape[0], maxind))
corres2 = -np.ones((FT.shape[0], maxind))
templength = 0
len_n = 0
## for source to target triangle coresspondence
rowlen = -1
for i in range(0, FS.shape[0]):
_, corresind = T_tree.query(VS_C[i, :],
k=maxind,
distance_upper_bound=thres)
corresind = corresind[corresind >= 0]
corresind = corresind[corresind < T_tree.data.shape[0]]
len_n = corresind.shape[0]
corresind[
np.sum(np.tile(NS[i, :], [NT[corresind, :].shape[0], 1]) *
NT[corresind, :],
axis=1) >= np.pi / 2] = -1
if len(corresind) != 0:
for j in range(0, len_n):
templength = np.max([
rowlen, corres2[corresind[j], :][
corres2[corresind[j], :] > -1].shape[0]
])
rowlen = corres2[corresind[j], :][
corres2[corresind[j], :] > -1].shape[0]
if rowlen == 10:
corres2[corresind[j], rowlen - 1] = i
else:
corres2[corresind[j], rowlen] = i
corres2 = corres2[:, 0:templength]
for i in range(0, FT.shape[0]):
_, corresind = S_tree.query(VT_C[i, :],
k=maxind,
distance_upper_bound=thres)
corresind = corresind[corresind >= 0]
corresind = corresind[corresind < S_tree.data.shape[0]]
templength = np.max([len_n, corresind.shape[0]])
len_n = corresind.shape[0]
corresind[
np.sum(np.tile(NT[i, :], [NS[corresind, :].shape[0], 1]) *
NS[corresind, :],
axis=1) >= np.pi / 2] = -1
corres1[i, 0:len_n] = corresind[0:len_n]
corres1 = corres1[:, 0:templength]
tempcorres = np.hstack((corres1, corres2))
corres = []
for i in range(0, FT.shape[0]):
temp = np.unique(tempcorres[i, :])
temp = temp[temp >= 0] # here delete -1 term
corres.append(temp)
save_pickle_file(FileNameTo_Save, corres)
return corres
class Test_sparse_distance_matrix_compiled(sparse_distance_matrix_consistency):
def setup_method(self):
n = 50
m = 4
np.random.seed(0)
data1 = np.random.randn(n, m)
data2 = np.random.randn(n, m)
self.T1 = cKDTree(data1, leafsize=2)
self.T2 = cKDTree(data2, leafsize=2)
self.ref_T1 = KDTree(data1, leafsize=2)
self.ref_T2 = KDTree(data2, leafsize=2)
self.r = 0.5
self.n = n
self.m = m
self.data1 = data1
self.data2 = data2
self.p = 2
def test_consistency_with_python(self):
M1 = self.T1.sparse_distance_matrix(self.T2, self.r)
M2 = self.ref_T1.sparse_distance_matrix(self.ref_T2, self.r)
assert_array_almost_equal(M1.todense(), M2.todense(), decimal=14)
def test_against_logic_error_regression(self):
# regression test for gh-5077 logic error
np.random.seed(0)
too_many = np.array(np.random.randn(18, 2), dtype=int)
tree = cKDTree(too_many, balanced_tree=False, compact_nodes=False)
d = tree.sparse_distance_matrix(tree, 3).todense()
assert_array_almost_equal(d, d.T, decimal=14)
def test_ckdtree_return_types(self):
# brute-force reference
ref = np.zeros((self.n, self.n))
for i in range(self.n):
for j in range(self.n):
v = self.data1[i, :] - self.data2[j, :]
ref[i, j] = np.dot(v, v)
ref = np.sqrt(ref)
ref[ref > self.r] = 0.
# test return type 'dict'
dist = np.zeros((self.n, self.n))
r = self.T1.sparse_distance_matrix(self.T2, self.r, output_type='dict')
for i, j in r.keys():
dist[i, j] = r[(i, j)]
assert_array_almost_equal(ref, dist, decimal=14)
# test return type 'ndarray'
dist = np.zeros((self.n, self.n))
r = self.T1.sparse_distance_matrix(self.T2,
self.r,
output_type='ndarray')
for k in range(r.shape[0]):
i = r['i'][k]
j = r['j'][k]
v = r['v'][k]
dist[i, j] = v
assert_array_almost_equal(ref, dist, decimal=14)
# test return type 'dok_matrix'
r = self.T1.sparse_distance_matrix(self.T2,
self.r,
output_type='dok_matrix')
assert_array_almost_equal(ref, r.todense(), decimal=14)
# test return type 'coo_matrix'
r = self.T1.sparse_distance_matrix(self.T2,
self.r,
output_type='coo_matrix')
assert_array_almost_equal(ref, r.todense(), decimal=14)
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
# Important: SET THIS TO FALSE IN PRODUCTION
self.simulate_lights = False
self.pose = None
self.waypoints = None
self.camera_image = None
self.lights = []
self.waypoint_tree = None
self.waypoints_received = False
# Initialize the classifier before subscribing to topics
if not self.simulate_lights:
rospy.loginfo('Loading Classifier')
self.light_classifier = TLClassifier()
rospy.loginfo('Loaded Classifier')
else:
rospy.logwarn('Simulated Traffic 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.listener = tf.TransformListener()
self.state = TrafficLight.UNKNOWN
self.last_state = TrafficLight.UNKNOWN
self.last_wp = -1
self.state_count = 0
self.loop()
def pose_cb(self, msg):
self.pose = msg
def waypoints_cb(self, lane):
self.waypoints = lane.waypoints
xy_waypoint_list = [[
waypoint.pose.pose.position.x, waypoint.pose.pose.position.y
] for waypoint in self.waypoints]
self.waypoint_tree = KDTree(xy_waypoint_list)
self.waypoints_received = True
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
def loop(self):
rate = rospy.Rate(2.) # 2.0Hz
while not rospy.is_shutdown():
if self.camera_image or self.simulate_lights:
self.publish_light_state()
rate.sleep()
def publish_light_state(self):
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, pose=None, x=None, y=None):
"""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
"""
if pose is not None:
_, nearest_waypoint = self.waypoint_tree.query(
[pose.position.x, pose.position.y])
else:
_, nearest_waypoint = self.waypoint_tree.query([x, y])
return nearest_waypoint
def get_simulated_light_state(self, light):
"""
Returns:
state: The state of the light from the simulator, which is always correct
"""
return light.state
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")
self.camera_image = None
#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
if not self.waypoints_received: return -1, TrafficLight.UNKNOWN
# 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)
else:
return -1, TrafficLight.UNKNOWN
#TODO find the closest visible traffic light (if one exists)
stop_line_indices = [
self.get_closest_waypoint(x=slp[0], y=slp[1])
for slp in stop_line_positions
]
nearest_stopline_dist = None
nearest_light_idx = None
for idx, sli in enumerate(stop_line_indices):
#TODO this is incomplete and does not account for wrapping:
distance = sli - car_position
if distance > 0 and nearest_stopline_dist is None or nearest_stopline_dist > distance:
nearest_stopline_dist = distance
nearest_light_idx = idx
if nearest_light_idx is not None and self.lights:
light = self.lights[nearest_light_idx]
light_wp = stop_line_indices[nearest_light_idx]
if light:
if self.simulate_lights:
state = self.get_simulated_light_state(light)
else:
state = self.get_light_state(light)
rospy.loginfo("classified image " + str(state))
return light_wp, state
return -1, TrafficLight.UNKNOWN
def setup_method(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_query_pairs_single_node():
tree = KDTree([[0, 1]])
assert_equal(tree.query_pairs(0.5), set())
def __init__(self, conf, problem, **kwargs):
from sfepy.discrete.state import State
from sfepy.discrete import Problem
from sfepy.base.conf import ProblemConf, get_standard_keywords
from scipy.spatial import cKDTree as KDTree
ScipyDirect.__init__(self, conf, **kwargs)
# init subproblems
pb_vars = problem.get_variables()
# get "master" DofInfo and last index
pb_adi_indx = problem.equations.variables.adi.indx
self.adi_indx = pb_adi_indx.copy()
last_indx = -1
for ii in six.itervalues(self.adi_indx):
last_indx = nm.max([last_indx, ii.stop])
# coupling variables
self.cvars_to_pb = {}
for jj in conf.coupling_variables:
self.cvars_to_pb[jj] = [None, None]
if jj in pb_vars.names:
if pb_vars[jj].dual_var_name is not None:
self.cvars_to_pb[jj][0] = -1
else:
self.cvars_to_pb[jj][1] = -1
# init subproblems
self.subpb = []
required, other = get_standard_keywords()
master_prefix = output.get_output_prefix()
for ii, ifname in enumerate(conf.others):
sub_prefix = master_prefix[:-1] + '-sub%d:' % (ii + 1)
output.set_output_prefix(sub_prefix)
kwargs['master_problem'] = problem
confi = ProblemConf.from_file(ifname,
required,
other,
define_args=kwargs)
pbi = Problem.from_conf(confi, init_equations=True)
sti = State(pbi.equations.variables)
pbi.equations.set_data(None, ignore_unknown=True)
pbi.time_update()
pbi.update_materials()
sti.apply_ebc()
pbi_vars = pbi.get_variables()
output.set_output_prefix(master_prefix)
self.subpb.append([pbi, sti, None])
# append "slave" DofInfo
for jj in pbi_vars.names:
if not (pbi_vars[jj].is_state()):
continue
didx = pbi.equations.variables.adi.indx[jj]
ndof = didx.stop - didx.start
if jj in self.adi_indx:
if ndof != \
(self.adi_indx[jj].stop - self.adi_indx[jj].start):
raise ValueError('DOFs do not match!')
else:
self.adi_indx.update(
{jj: slice(last_indx, last_indx + ndof, None)})
last_indx += ndof
for jj in conf.coupling_variables:
if jj in pbi_vars.names:
if pbi_vars[jj].dual_var_name is not None:
self.cvars_to_pb[jj][0] = ii
else:
self.cvars_to_pb[jj][1] = ii
self.subpb.append([problem, None, None])
self.cvars_to_pb_map = {}
for varname, pbs in six.iteritems(self.cvars_to_pb):
# match field nodes
coors = []
for ii in pbs:
pbi = self.subpb[ii][0]
pbi_vars = pbi.get_variables()
fcoors = pbi_vars[varname].field.coors
dc = nm.abs(nm.max(fcoors, axis=0)\
- nm.min(fcoors, axis=0))
ax = nm.where(dc > 1e-9)[0]
coors.append(fcoors[:, ax])
if len(coors[0]) != len(coors[1]):
raise ValueError('number of nodes does not match!')
kdtree = KDTree(coors[0])
map_12 = kdtree.query(coors[1])[1]
pbi1 = self.subpb[pbs[0]][0]
pbi1_vars = pbi1.get_variables()
eq_map_1 = pbi1_vars[varname].eq_map
pbi2 = self.subpb[pbs[1]][0]
pbi2_vars = pbi2.get_variables()
eq_map_2 = pbi2_vars[varname].eq_map
dpn = eq_map_2.dpn
nnd = map_12.shape[0]
map_12_nd = nm.zeros((nnd * dpn, ), dtype=nm.int32)
if dpn > 1:
for ii in range(dpn):
map_12_nd[ii::dpn] = map_12 * dpn + ii
else:
map_12_nd = map_12
idx = nm.where(eq_map_2.eq >= 0)[0]
self.cvars_to_pb_map[varname] = eq_map_1.eq[map_12[idx]]
def setup_method(self):
n = 50
m = 2
np.random.seed(1234)
self.T1 = KDTree(np.random.randn(n, m), leafsize=2)
self.T2 = KDTree(np.random.randn(n, m), leafsize=2)
def find_nearest_node(self, x, y, node):
tree = KDTree(np.vstack((x, y)).T)
distance, index = tree.query(np.array([node.x, node.y]), eps=0.5)
return index
def setup_method(self):
Test_small.setup_method(self)
self.kdtree = KDTree(self.data, leafsize=1)
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.waypoint_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, Loader=yaml.FullLoader)
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, lane):
self.waypoints = lane
if not self.waypoints_2d:
self.waypoints_2d = [[waypoint.pose.pose.position.x, waypoint.pose.pose.position.y] for waypoint in
lane.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
"""
# Only process every 10th image
self.has_image = True
self.camera_image = msg
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 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)
"""
# TODO Remove... For testing, just return the 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_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 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 d >= 0 and d < diff:
diff = d
closest_light = light
line_wp_idx = temp_wp_idx
else:
rospy.logerr("pose not set")
if closest_light:
state = self.get_light_state(light)
return line_wp_idx, state
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 = []
self.bridge = CvBridge()
self.light_classifier = TLClassifier(rospy.get_param('~model_file'))
#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.class_count = 0
self.process_count = 0
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
rospy.Subscriber('/current_pose',
PoseStamped,
self.pose_cb,
queue_size=2)
rospy.Subscriber('/base_waypoints',
Lane,
self.waypoints_cb,
queue_size=8)
rospy.Subscriber('/vehicle/traffic_lights',
TrafficLightArray,
self.traffic_cb,
queue_size=2)
rospy.Subscriber('/image_color', Image, self.image_cb, queue_size=1)
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
rospy.spin()
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
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
:param msg: image from car-mounted camera
"""
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 process_traffic_lights(self):
"""
Finds closest visible traffic light, if one exists, and determines its location and color
:return:
int: index of waypoint closest 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 = -1
state = TrafficLight.UNKNOWN
# 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)
diff = len(self.base_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 we have found a closest light to monitor, then determine the stop line position of this light
if closest_light:
self.process_count += 1
state = self.get_light_state(closest_light)
if (self.process_count % LOGGING_THROTTLE_FACTOR) == 0:
rospy.logwarn("DETECT: line_wp_idx={}, state={}".format(
line_wp_idx, self.to_string(state)))
return line_wp_idx, state
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
:param x: x coord position to match a waypoint to
:param y: y coord position to match a waypoint to
:return: index of the closest waypoint in self.waypoints
"""
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
:param light: light to classify
:return: ID of traffic light color (specified in styx_msgs/TrafficLight)
"""
# For test mode, just return the light state
if TEST_MODE_ENABLED:
classification = light.state
else:
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
# Get classification
classification = self.light_classifier.get_classification(cv_image)
# Save image (throttled)
if SAVE_IMAGES and (self.process_count % LOGGING_THROTTLE_FACTOR
== 0):
save_file = "../../../imgs/{}-{:.0f}.jpeg".format(
self.to_string(classification), (time.time() * 100))
cv2.imwrite(save_file, cv_image)
return classification
def to_string(self, state):
out = "unknown"
if state == TrafficLight.GREEN:
out = "green"
elif state == TrafficLight.YELLOW:
out = "yellow"
elif state == TrafficLight.RED:
out = "red"
return out
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
# Parameters
config_string = rospy.get_param("/traffic_light_config")
self.config = yaml.load(config_string)
# Subscribers
self.current_pose_sub = rospy.Subscriber('/current_pose', PoseStamped,
self.pose_cb)
self.base_waypoints_sub = rospy.Subscriber('/base_waypoints', Lane,
self.base_waypoints_cb)
self.traffic_lights_sub = rospy.Subscriber('/vehicle/traffic_lights',
TrafficLightArray,
self.traffic_cb)
self.image_color_sub = rospy.Subscriber(
'/image_color', Image,
self.image_cb) #TODO: evaluate using image_raw instead
# Publishers
self.upcoming_red_light_pub = rospy.Publisher('/traffic_waypoint',
Int32,
queue_size=1)
# TL detector node variables
self.current_pose = None
self.base_waypoints = None
self.base_waypoints_init = False
self.base_waypoints_tree = None
self.camera_image = None
self.current_traffic_lights = []
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.step()
def step(self):
rate = rospy.Rate(50)
while not rospy.is_shutdown():
rate.sleep()
'''
This method updates the current ego vehicle pose
'''
def pose_cb(self, data):
self.current_pose = data
'''
This method stores the base waypoints in the class variables
'''
def base_waypoints_cb(self, input_waypoints):
if not self.base_waypoints_init:
self.base_waypoints = input_waypoints
self.base_waypoints_2d = [[
wp.pose.pose.position.x, wp.pose.pose.position.y
] for wp in input_waypoints.waypoints]
self.base_waypoints_tree = KDTree(self.base_waypoints_2d)
self.base_waypoints_init = True
#rospy.loginfo('Base waypoints initialized.')
'''
This method stores the traffic light data in the class variables
'''
def traffic_cb(self, data):
self.current_traffic_lights = data.lights
'''
This method identifies the current state of the traffic light in the incoming front camera sensor image
'''
def image_cb(self, data):
self.has_image = True
self.camera_image = data
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
red = state == TrafficLight.RED or state == TrafficLight.UNKNOWN
light_wp = light_wp if 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
'''
This method identifies the closest path waypoint to a given position
'''
def get_closest_waypoint(self, pose_x, pose_y):
closest_wp_idx = self.base_waypoints_tree.query([pose_x, pose_y], 1)[1]
return closest_wp_idx
'''
This method determines the current state of the traffic light
'''
def get_light_state(self, light):
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)
'''
This method finds the closest traffic light within a max distance and determines the distance to its stop line and its state
'''
def process_traffic_lights(self):
closest_line_wp_idx = None
closest_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']
if (self.current_pose):
current_pose_wp_idx = self.get_closest_waypoint(
self.current_pose.pose.position.x,
self.current_pose.pose.position.y)
diff = len(self.base_waypoints.waypoints)
for i, current_light in enumerate(self.current_traffic_lights):
line = stop_line_positions[i]
current_line_wp_idx = self.get_closest_waypoint(
line[0], line[1])
d = current_line_wp_idx - current_pose_wp_idx
if d >= 0 and d < diff and d < MAX_STOP_LINE_DIST:
diff = d
closest_light = current_light
closest_line_wp_idx = current_line_wp_idx
if closest_light:
#rospy.loginfo('Close traffic light found.')
closest_light_state = self.get_light_state(closest_light)
#closest_light_state = closest_light.state
#rospy.loginfo('Current pose waypoint index: x = %s', current_pose_wp_idx)
#rospy.loginfo('Closest stop line waypoint index: x = %s', closest_line_wp_idx)
#rospy.loginfo('Closest traffic light state: %s', closest_light_state)
if closest_light_state is None:
closest_light_state = TrafficLight.UNKNOWN
#if closest_light_state == TrafficLight.UNKNOWN:
#rospy.loginfo('Unknown Traffic light state.')
return closest_line_wp_idx, closest_light_state
#rospy.loginfo('No close traffic light detected.')
return -1, TrafficLight.UNKNOWN
def waypoints_cb(self, lane):
self.waypoints = lane
if not self.waypoints_2d:
self.waypoints_2d = [[waypoint.pose.pose.position.x, waypoint.pose.pose.position.y] for waypoint in
lane.waypoints]
self.waypoint_tree = KDTree(self.waypoints_2d)
return int(h[0:4], 16), int(('0x' + h[4:6]), 16), int(('0x' + h[6:8]), 16)
filename = input("What's the image name? ")
new_w, new_h = map(
int,
raw_input("What's the new height x width? Like 28 28. ").split(' '))
palette_hex = [
'0xFF0000', '0xF83800', '0xF0D0B0', '0x503000', '0xFFE0A8', '0x0058F8',
'0xFCFCFC', '0xBCBCBC', '0xA40000', '0xD82800', '0xFC7460', '0xFCBCB0',
'0xF0BC3C', '0xAEACAE', '0x363301', '0x6C6C01', '0xBBBD00', '0x88D500',
'0x398802', '0x65B0FF', '0x155ED8', '0x800080', '0x24188A'
]
palette_rgb = [hex_to_rgb(color) for color in palette_hex]
pixel_tree = KDTree(palette_rgb)
im = Image.open("./sprite_originals/" + filename +
".png") #Can be many different formats.
im = im.convert("RGBA")
im = im.resize((new_w, new_h), Image.ANTIALIAS) # regular resize
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("./sprite_bytes/" + filename + '.txt', 'w')
i = 0
for y in range(im.size[1]):
for x in range(im.size[0]):
def findfiducials(spots, input_transform=None, separation=8.):
global metrology_pinholes_table
global metrology_fiducials_table
log = get_logger()
log.debug(
"load input tranformation we will use to go from FP to FVC pixels")
if input_transform is None:
input_transform = resource_filename('desimeter',
"data/single-lens-fvc2fp.json")
log.info("loading input tranform from {}".format(input_transform))
try:
input_tx = FVCFP_ZhaoBurge.read_jsonfile(input_transform)
except AssertionError as e:
log.warning(
"Failed to read input transfo as Zhao Burge, try polynomial...")
input_tx = FVCFP_Polynomial.read_jsonfile(input_transform)
if metrology_pinholes_table is None:
filename = resource_filename('desimeter', "data/fp-metrology.csv")
if not os.path.isfile(filename):
log.error("cannot find {}".format(filename))
raise IOError("cannot find {}".format(filename))
log.info("reading metrology in {}".format(filename))
metrology_table = Table.read(filename, format="csv")
log.debug("keep only the pinholes")
metrology_pinholes_table = metrology_table[:][
metrology_table["PINHOLE_ID"] > 0]
# use input transform to convert X_FP,Y_FP to XPIX,YPIX
xpix, ypix = input_tx.fp2fvc(metrology_pinholes_table["X_FP"],
metrology_pinholes_table["Y_FP"])
metrology_pinholes_table["XPIX"] = xpix
metrology_pinholes_table["YPIX"] = ypix
log.debug("define fiducial location as central dot")
metrology_fiducials_table = metrology_pinholes_table[:][
metrology_pinholes_table["PINHOLE_ID"] == 4]
# find fiducials candidates
log.info("select spots with at least two close neighbors (in pixel units)")
xy = np.array([spots["XPIX"], spots["YPIX"]]).T
tree = KDTree(xy)
measured_spots_distances, measured_spots_indices = tree.query(
xy, k=4, distance_upper_bound=separation)
number_of_neighbors = np.sum(measured_spots_distances < separation, axis=1)
fiducials_candidates_indices = np.where(
number_of_neighbors >= 3)[0] # including self, so at least 3 pinholes
# match candidates to fiducials from metrology
log.info(
"first match {} fiducials candidates to metrology ({}) with iterative fit"
.format(fiducials_candidates_indices.size,
len(metrology_fiducials_table)))
x1 = spots["XPIX"][fiducials_candidates_indices]
y1 = spots["YPIX"][fiducials_candidates_indices]
x2 = metrology_fiducials_table["XPIX"] # do I need to do this?
y2 = metrology_fiducials_table["YPIX"]
nloop = 20
saved_median_distance = 0
for loop in range(nloop):
indices_2, distances = match_same_system(x1, y1, x2, y2)
mdist = np.median(distances[indices_2 >= 0])
if loop < nloop - 1:
maxdistance = max(10, 3. * 1.4 * mdist)
else: # final iteration
maxdistance = 10 # pixel
selection = np.where((indices_2 >= 0) & (distances < maxdistance))[0]
log.info("iter #{} median_dist={} max_dist={} matches={}".format(
loop, mdist, maxdistance, selection.size))
corr21 = SimpleCorr()
corr21.fit(x2[indices_2[selection]], y2[indices_2[selection]],
x1[selection], y1[selection])
x2, y2 = corr21.apply(x2, y2)
if np.abs(saved_median_distance - mdist) < 0.0001:
break # no more improvement
saved_median_distance = mdist
# use same coord system match (note we now match the otherway around)
indices_1, distances = match_same_system(x2, y2, x1, y1)
maxdistance = 10. # FVC pixels
selection = np.where((indices_1 >= 0) & (distances < maxdistance))[0]
fiducials_candidates_indices = fiducials_candidates_indices[
indices_1[selection]]
matching_known_fiducials_indices = selection
log.debug(
"mean distance = {:4.2f} pixels for {} matched and {} known fiducials".
format(np.mean(distances[distances < maxdistance]),
fiducials_candidates_indices.size,
metrology_fiducials_table["XPIX"].size))
log.debug("now matching pinholes ...")
nspots = spots["XPIX"].size
if 'LOCATION' not in spots.dtype.names:
spots.add_column(Column(np.zeros(nspots, dtype=int)), name='LOCATION')
if 'PINHOLE_ID' not in spots.dtype.names:
spots.add_column(Column(np.zeros(nspots, dtype=int)),
name='PINHOLE_ID')
for index1, index2 in zip(fiducials_candidates_indices,
matching_known_fiducials_indices):
location = metrology_fiducials_table["LOCATION"][index2]
# get indices of all pinholes for this matched fiducial
# note we now use the full pinholes metrology table
pi1 = measured_spots_indices[index1][
measured_spots_distances[index1] < separation]
pi2 = np.where(metrology_pinholes_table["LOCATION"] == location)[0]
x1 = spots["XPIX"][pi1]
y1 = spots["YPIX"][pi1]
x2 = metrology_pinholes_table["XPIX"][pi2]
y2 = metrology_pinholes_table["YPIX"][pi2]
indices_2, distances = match_arbitrary_translation_dilatation(
x1, y1, x2, y2)
metrology_pinhole_ids = metrology_pinholes_table["PINHOLE_ID"][pi2]
pinhole_ids = np.zeros(x1.size, dtype=int)
matched = (indices_2 >= 0)
pinhole_ids[matched] = metrology_pinhole_ids[indices_2[matched]]
spots["LOCATION"][pi1[matched]] = location
spots["PINHOLE_ID"][pi1[matched]] = pinhole_ids[matched]
if np.sum(pinhole_ids == 0) > 0:
log.warning(
"only matched pinholes {} for {} detected at LOCATION {} xpix~{} ypix~{}"
.format(pinhole_ids[pinhole_ids > 0], x1.size, location,
int(np.mean(x1)), int(np.mean(y1))))
# check duplicates
if np.unique(
pinhole_ids[pinhole_ids > 0]).size != np.sum(pinhole_ids > 0):
xfp = np.mean(metrology_pinholes_table[pi2]["X_FP"])
yfp = np.mean(metrology_pinholes_table[pi2]["Y_FP"])
log.warning(
"duplicate(s) pinhole ids in {} at LOCATION={} xpix~{} ypix~{} xfp~{} yfp~{}"
.format(pinhole_ids, location, int(np.mean(x1)),
int(np.mean(y1)), int(xfp), int(yfp)))
bc = np.bincount(pinhole_ids[pinhole_ids > 0])
duplicates = np.where(bc > 1)[0]
for duplicate in duplicates:
log.warning(
"Unmatch ambiguous pinhole id = {}".format(duplicate))
selection = (spots["LOCATION"]
== location) & (spots["PINHOLE_ID"] == duplicate)
spots["PINHOLE_ID"][selection] = 0
spots["PETAL_LOC"] = spots["LOCATION"] // 1000
spots["DEVICE_LOC"] = spots["LOCATION"] % 1000
n_matched_pinholes = np.sum(spots["PINHOLE_ID"] > 0)
n_matched_fiducials = np.sum(spots["PINHOLE_ID"] == 4)
log.info("matched {} pinholes from {} fiducials".format(
n_matched_pinholes, n_matched_fiducials))
return spots
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)
class TLDetector(object):
def __init__(self):
rospy.init_node('tl_detector')
self.on_simulator = rospy.get_param('~on_simulator')
self.use_ground_truth = False
self.pose = None
self.waypoints = None
self.camera_image = None
self.camera_image_count = 1
self.lights = []
self.waypoint_tree = None
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)
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.on_simulator)
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):
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)
#rospy.logwarn(self.waypoint_tree)
def traffic_cb(self, msg):
self.lights = msg.lights
#rospy.loginfo("tl_detector::traffic_cb: {0}".format(self.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
"""
# only process every n-th image to improve performance.
if self.camera_image_count < 10:
self.camera_image_count += 1
self.has_image = False
self.camera_image = None
else:
self.camera_image_count = 1
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
"""
#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)
"""
# For testing, just return the 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)
if self.use_ground_truth:
rospy.loginfo("debugging, using ground truth")
return light.state
if (not self.has_image):
# rospy.loginfo("no image info!")
self.prev_light_loc = None
return TrafficLight.RED
cv_image = self.bridge.imgmsg_to_cv2(self.camera_image, "bgr8")
##Get classification
# return self.light_classifier.get_classification(cv_image)
pred_state = self.light_classifier.get_classification(cv_image)
rospy.loginfo("SSD network says: %s (wright answer is %s)", pred_state,
light.state)
return pred_state
def euclidian_distance(self, position1, position2):
x, y, z = position1.x - position2.x, position1.y - position2.y, position1.z - position2.z
return math.sqrt(x * x + y * y + z * z)
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])
# Find 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
# check the distance from car to closest_light
closest_light_distance = self.euclidian_distance(
self.pose.pose.position,
self.waypoints.waypoints[line_wp_idx].pose.pose.position)
if closest_light and closest_light_distance < LIGHT_DISTANCE_THRESHOLD:
state = self.get_light_state(
closest_light
) # approaching a light, try to determine its state
rospy.loginfo("approaching %s traffic light %f ahead", state,
closest_light_distance)
if state == TrafficLight.RED:
return line_wp_idx, state
else:
return -1, TrafficLight.UNKNOWN
else: # far away from light, hence state is don't care.
return -1, TrafficLight.UNKNOWN