fix some bugs on 2d data

Author: pattonw

micron/network/nms.py

@@ -60,7 +60,7 @@ def max_detection(
 
     # not sure if this does anything
     conv_filter = np.ones([1 for _ in range(n_dim + 2)])
-    upsampled = tf.nn.conv3d_transpose(
+    upsampled = conv_transpose(
         max_pool,
         conv_filter.astype(np.float16),
         [1, *sm_dims, 1],
@@ -73,5 +73,8 @@ def max_detection(
     reduced_maxima = tf.equal(upsampled, sm_maxima)
     reduced_maxima = tf.logical_and(reduced_maxima, sm_maxima > 1)
 
-    return maxima[0, :, :, :, 0], reduced_maxima[0, :, :, :, 0]
+    if n_dim == 2:
+        return maxima[0, :, :, 0], reduced_maxima[0, :, :, 0]
+    elif n_dim == 3:
+        return maxima[0, :, :, :, 0], reduced_maxima[0, :, :, :, 0]