Skip to content

Commit 025f4c1

Browse files
julienrwaleedka
authored andcommitted
Fix problem with argmax on (0,0) arrays.
Fix matterport#170
1 parent d0c52d7 commit 025f4c1

File tree

1 file changed

+5
-1
lines changed

1 file changed

+5
-1
lines changed

mrcnn/model.py

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -560,7 +560,11 @@ def detection_targets_graph(proposals, gt_class_ids, gt_boxes, gt_masks, config)
560560

561561
# Assign positive ROIs to GT boxes.
562562
positive_overlaps = tf.gather(overlaps, positive_indices)
563-
roi_gt_box_assignment = tf.argmax(positive_overlaps, axis=1)
563+
roi_gt_box_assignment = tf.cond(
564+
tf.greater(tf.shape(positive_overlaps)[1], 0),
565+
true_fn = lambda: tf.argmax(positive_overlaps, axis=1),
566+
false_fn = lambda: tf.cast(tf.constant([]),tf.int64)
567+
)
564568
roi_gt_boxes = tf.gather(gt_boxes, roi_gt_box_assignment)
565569
roi_gt_class_ids = tf.gather(gt_class_ids, roi_gt_box_assignment)
566570

0 commit comments

Comments
 (0)