Fix key error bug in iou

funkelab · Apr 3, 2024 · 0f461cb · 0f461cb
1 parent 7820cca
commit 0f461cb
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/motile_toolbox/candidate_graph/iou.py b/src/motile_toolbox/candidate_graph/iou.py
@@ -104,12 +104,15 @@ def add_iou(
     if node_frame_dict is None:
         node_frame_dict = _compute_node_frame_dict(cand_graph)
     frames = sorted(node_frame_dict.keys())
+    print(frames)
     ious = _get_iou_dict(segmentation, multihypo=multihypo)
+    print(ious)
     for frame in tqdm(frames):
-        if frame + 1 not in node_frame_dict:
+        if frame + 1 not in node_frame_dict.keys():
             continue
         next_nodes = node_frame_dict[frame + 1]
         for node_id in node_frame_dict[frame]:
             for next_id in next_nodes:
                 iou = ious.get(node_id, {}).get(next_id, 0)
-                cand_graph.edges[(node_id, next_id)][EdgeAttr.IOU.value] = iou
+                if (node_id, next_id) in cand_graph.edges:
+                    cand_graph.edges[(node_id, next_id)][EdgeAttr.IOU.value] = iou