TFDetect dynamic anchor count assignment fix (#5668)

* fix tf.py when anchors not equal to 3

* revert the isort fix

* update the fix to use anchor attribute available already
This commit is contained in:
Nrupatunga authored and GitHub committed 2021-11-16 12:36:00 +01:00
1 parent e80a09bbfa
commit 0453b758e7
1 file changed
+1 -1
+1 -1
View File
@@ -233,7 +233,7 @@ class TFDetect(keras.layers.Layer):
xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
y = tf.concat([xy, wh, y[..., 4:]], -1)
z.append(tf.reshape(y, [-1, 3 * ny * nx, self.no]))
z.append(tf.reshape(y, [-1, self.na * ny * nx, self.no]))
return x if self.training else (tf.concat(z, 1), x)