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()