Nrupatunga commited on
Commit
0453b75
·
unverified ·
1 Parent(s): e80a09b

TFDetect dynamic anchor count assignment fix (#5668)

Browse files

* fix tf.py when anchors not equal to 3

* revert the isort fix

* update the fix to use anchor attribute available already

Files changed (1) hide show
  1. models/tf.py +1 -1
models/tf.py CHANGED
@@ -233,7 +233,7 @@ class TFDetect(keras.layers.Layer):
233
  xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
234
  wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
235
  y = tf.concat([xy, wh, y[..., 4:]], -1)
236
- z.append(tf.reshape(y, [-1, 3 * ny * nx, self.no]))
237
 
238
  return x if self.training else (tf.concat(z, 1), x)
239
 
 
233
  xy /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
234
  wh /= tf.constant([[self.imgsz[1], self.imgsz[0]]], dtype=tf.float32)
235
  y = tf.concat([xy, wh, y[..., 4:]], -1)
236
+ z.append(tf.reshape(y, [-1, self.na * ny * nx, self.no]))
237
 
238
  return x if self.training else (tf.concat(z, 1), x)
239