def test_ia_util_14():
"""
Some tests of our KDTree.
"""
x1 = numpy.array([1.0, 2.0, 3.0])
y1 = numpy.array([1.0, 1.0, 1.0])
kd = iaUtilsC.KDTree(x=x1, y=y1)
x2 = numpy.array([1.1, 4.0])
y2 = numpy.array([1.0, 1.0])
[dist, index] = kd.nearest(x2, y2, 0.2)
assert (abs(dist[0] - 0.1) < 1.0e-6)
assert (abs(dist[1] + 1.0) < 1.0e-6)
assert (index[0] == 0)
assert (index[1] == -1)
[dist, index] = kd.nearest(x2, y2, 3.1)
assert (abs(dist[0] - 0.1) < 1.0e-6)
assert (abs(dist[1] - 1.0) < 1.0e-6)
assert (index[0] == 0)
assert (index[1] == 2)
kd.cleanup()
def trackFiducials(sa_hdf5_filename, max_gap=0, radius=0.0, reference_frame=0):
"""
max_gap - The maximum number of frames with no objects before a track is
considered to be terminated.
radius - Maximum distance for an object to be in a track in pixels.
reference_frame - The localizations in this frame will be used as the
fiducial reference locations.
"""
with SAH5Fiducials(sa_hdf5_filename) as h5:
ref_locs = h5.getLocalizationsInFrame(reference_frame,
fields=["x", "y"])
if not ref_locs:
raise FiducialException("No localizations in frame " +
str(reference_frame))
# Create fiducials list.
fiducials = []
for i in range(ref_locs["x"].size):
fd = Fiducial(track_id=i)
fd.addLocalization(ref_locs, i)
fiducials.append(fd)
h5.setNFiducials(len(fiducials))
# Iterate over localizations.
for fnum, locs in h5.localizationsIterator(skip_empty=False,
fields=["x", "y"]):
# User feedback.
if ((fnum % 500) == 0):
print(" processing frame {0:0d}, {1:0d} fiducials".format(
fnum, len(fiducials)))
# Check if we still have any 'live' fiducials.
if (len(fiducials) == 0):
# Mark all localizations as non-fiducial, i.e. -1.
if bool(locs):
locs_fd_id = numpy.zeros(locs["x"].size,
dtype=numpy.int32) - 1
h5.addFiducialID(locs_fd_id, fnum)
continue
# Check that the frame had localizations, assign them to fiducials if it did.
if bool(locs):
# Create numpy array for storage of the fiducial id for each localization.
locs_fd_id = numpy.zeros(locs["x"].size, dtype=numpy.int32) - 1
# Create arrays with current fiducial centers. This is also increments
# the fiducials last added counter.
tx = numpy.zeros(len(fiducials))
ty = numpy.zeros(len(fiducials))
for i, elt in enumerate(fiducials):
[tx[i], ty[i]] = elt.getCenter()
elt.incLastAdded()
kd_locs = iaUtilsC.KDTree(locs["x"], locs["y"])
# Query for nearest localization to each fiducial.
index_fd = kd_locs.nearest(tx, ty, radius)[1]
# Add fiducial information to localizations.
for i in range(index_fd.size):
if (index_fd[i] > -1):
fiducials[i].addLocalization(locs, index_fd[i])
locs_fd_id[index_fd[i]] = fiducials[i].getFiducialId()
h5.addFiducialID(locs_fd_id, fnum)
# Clean up KD trees.
kd_locs.cleanup()
# Otherwise just increment the fiducials last added counter.
else:
for elt in fiducials:
elt.incLastAdded()
# Remove fiducials that have not had any localizations added for max_gap frames.
temp = fiducials
fiducials = []
for elt in temp:
if (elt.getLastAdded() <= max_gap):
fiducials.append(elt)
else:
print("Lost fiducial", elt.getFiducialId(), "at frame",
fnum)
def tracker(sa_hdf5_filename, descriptor="", max_gap=0, radius=0.0):
"""
descriptor - A string containing the frame designation.
max_gap - The maximum number of frames with no objects before a track is
considered to be terminated.
radius - Maximum distance for an object to be in a track in pixels.
"""
# Just set localization category if radius is zero or negative.
if (radius <= 0.0):
with saH5Py.SAH5Py(sa_hdf5_filename) as h5:
for fnum, locs in h5.localizationsIterator():
# Determine current frame description.
fdesc = 1
if (len(descriptor) > 0):
fdesc = int(descriptor[(fnum % len(descriptor))])
# The category is the descriptor minus 1.
category = fdesc - 1
h5.addCategory(category, fnum)
# Otherwise do the tracking.
else:
track_id = 0
current_tracks = []
with saH5Py.SAH5Py(sa_hdf5_filename) as h5:
tw = TrackWriter(h5)
for fnum, locs in h5.localizationsIterator(skip_empty=False):
# User feedback.
if ((fnum % 500) == 0):
print(" processing frame {0:0d}, {1:0d} tracks".format(
fnum, track_id))
# Determine current frame description.
fdesc = 1
if (len(descriptor) > 0):
fdesc = int(descriptor[(fnum % len(descriptor))])
# The category is the descriptor minus 1.
category = fdesc - 1
# Add/update localization category.
if bool(locs):
h5.addCategory(category, fnum)
# Go to the next frame if this is an activation frame.
if (fdesc == 0):
continue
# Check that the frame had localizations, assign them if it did.
index_locs = None
locs_track_id = None
if bool(locs):
# Create numpy array for storage of the track id for each localization.
locs_track_id = numpy.zeros(locs["x"].size,
dtype=numpy.int64)
# Create arrays with current track centers. This is also increments
# the tracks last added counter.
tx = numpy.zeros(len(current_tracks))
ty = numpy.zeros(len(current_tracks))
for i, elt in enumerate(current_tracks):
[tx[i], ty[i]] = elt.getCenter()
elt.incLastAdded()
kd_locs = iaUtilsC.KDTree(locs["x"], locs["y"])
kd_tracks = iaUtilsC.KDTree(tx, ty)
# Query KD trees.
index_locs = kd_tracks.nearest(locs["x"], locs["y"],
radius)[1]
index_tracks = kd_locs.nearest(tx, ty, radius)[1]
# Add localizations to tracks. The localization must be the closest
# one to the track and vice-versa. We're trying to avoid multiple
# localizations in a single frame in the track, and one localization
# in multiple tracks.
#
for i in range(locs["x"].size):
if (index_locs[i] > -1):
# Check that the track and the localization agree that each
# is closest to the other.
#
tr = None
if (index_tracks[index_locs[i]] == i):
tr = current_tracks[index_locs[i]]
tr.addLocalization(locs, i)
else:
tr = Track(category=category,
frame_number=fnum,
track_id=track_id)
tr.addLocalization(locs, i)
current_tracks.append(tr)
track_id += 1
locs_track_id[i] = tr.track_id
# Clean up KD trees.
kd_locs.cleanup()
kd_tracks.cleanup()
# Otherwise just increment the current tracks last added counter.
else:
for elt in current_tracks:
elt.incLastAdded()
# Remove tracks that have not had any localizations added for
# max_gap frames.
temp = current_tracks
current_tracks = []
for elt in temp:
if (elt.getLastAdded() > max_gap):
tw.writeTrack(elt)
else:
current_tracks.append(elt)
# Start new tracks from the localizations that were not in
# a track.
if index_locs is not None:
for i in range(locs["x"].size):
if (index_locs[i] < 0):
tr = Track(category=category,
frame_number=fnum,
track_id=track_id)
tr.addLocalization(locs, i)
current_tracks.append(tr)
locs_track_id[i] = tr.track_id
track_id += 1
# Add track information for localizations.
if locs_track_id is not None:
h5.addTrackID(locs_track_id, fnum)
# Write the remaining tracks & close the track writer.
for elt in current_tracks:
tw.writeTrack(elt)
tw.finish()