Spaces:
Sleeping
Sleeping
File size: 5,143 Bytes
aff3c6f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
import logging
from itertools import chain
import networkx as nx
import numpy as np
import scipy
from tqdm import tqdm
from .track_graph import TrackGraph
from typing import Optional, Tuple
# from trackastra.tracking import graph_to_napari_tracks, graph_to_ctc
logger = logging.getLogger(__name__)
def copy_edge(edge: tuple, source: nx.DiGraph, target: nx.DiGraph):
if edge[0] not in target.nodes:
target.add_node(edge[0], **source.nodes[edge[0]])
if edge[1] not in target.nodes:
target.add_node(edge[1], **source.nodes[edge[1]])
target.add_edge(edge[0], edge[1], **source.edges[(edge[0], edge[1])])
def track_greedy(
candidate_graph,
allow_divisions=True,
threshold=0.5,
edge_attr="weight",
):
"""Greedy matching, global.
Iterates over global edges sorted by weight, and keeps edge if feasible and weight above threshold.
Args:
allow_divisions (bool, optional):
Whether to model divisions. Defaults to True.
Returns:
solution_graph: NetworkX graph of tracks
"""
logger.info("Running greedy tracker")
solution_graph = nx.DiGraph()
# TODO bring back
# if args.gt_as_dets:
# solution_graph.add_nodes_from(candidate_graph.nodes(data=True))
edges = candidate_graph.edges(data=True)
edges = sorted(
edges,
key=lambda edge: edge[2][edge_attr],
reverse=True,
)
for edge in tqdm(edges, desc="Greedily matched edges"):
node_in, node_out, features = edge
assert (
features[edge_attr] <= 1.0
), "Edge weights are assumed to be normalized to [0,1]"
# assumes sorted edges
if features[edge_attr] < threshold:
break
# Check whether this edge is a feasible edge to add
# i.e. no fusing
if node_out in solution_graph.nodes and solution_graph.in_degree(node_out) > 0:
# target node already has an incoming edge
continue
if node_in in solution_graph and solution_graph.out_degree(node_in) >= (
2 if allow_divisions else 1
):
# parent node already has max number of outgoing edges
continue
# otherwise add to solution
copy_edge(edge, candidate_graph, solution_graph)
# df, masks = graph_to_ctc(solution_graph, masks_original)
# tracks, tracks_graph, _ = graph_to_napari_tracks(solution_graph)
return solution_graph
# TODO this should all be in a tracker class
# return df, masks, solution_graph, tracks_graph, tracks, candidate_graph
def build_graph(
nodes: dict,
weights: Optional[tuple] = None,
use_distance: bool = False,
max_distance: Optional[int] = None,
max_neighbors: Optional[int] = None,
delta_t=1,
):
logger.info(f"Build candidate graph with {delta_t=}")
G = nx.DiGraph()
for node in nodes:
G.add_node(
node["id"],
time=node["time"],
label=node["label"],
coords=node["coords"],
# index=node["index"],
weight=1,
)
if use_distance:
weights = None
if weights:
weights = {w[0]: w[1] for w in weights}
graph = TrackGraph(G, frame_attribute="time")
frame_pairs = zip(
chain(*[
list(range(graph.t_begin, graph.t_end - d)) for d in range(1, delta_t + 1)
]),
chain(*[
list(range(graph.t_begin + d, graph.t_end)) for d in range(1, delta_t + 1)
]),
)
iterator = tqdm(
frame_pairs,
total=(graph.t_end - graph.t_begin) * delta_t,
leave=False,
)
for t_begin, t_end in iterator:
n_edges_t = len(G.edges)
ni, nj = graph.nodes_by_frame(t_begin), graph.nodes_by_frame(t_end)
pi = []
for _ni in ni:
pi.append(np.array(G.nodes[_ni]["coords"]))
pi = np.stack(pi)
pj = []
for _nj in nj:
pj.append(np.array(G.nodes[_nj]["coords"]))
pj = np.stack(pj)
dists = scipy.spatial.distance.cdist(pi, pj)
for _i, _ni in enumerate(ni):
inds = np.argsort(dists[_i])
neighbors = 0
for _j, _nj in zip(inds, np.array(nj)[inds]):
if max_neighbors and neighbors >= max_neighbors:
break
dist = dists[_i, _j]
if max_distance is None or dist <= max_distance:
if weights is None:
G.add_edge(_ni, _nj, weight=1 - dist / max_distance)
neighbors += 1
else:
if (_ni, _nj) in weights:
G.add_edge(_ni, _nj, weight=weights[(_ni, _nj)])
neighbors += 1
e_added = len(G.edges) - n_edges_t
if e_added == 0:
logger.warning(f"No candidate edges in frame {t_begin}")
iterator.set_description(
f"{e_added} edges in frame {t_begin} Total edges: {len(G.edges)}"
)
logger.info(f"Added {len(G.nodes)} vertices, {len(G.edges)} edges")
return G
|