update

Longer-Liushen · Feb 15, 2020 · 2d0e73f · 2d0e73f
1 parent ffea584
commit 2d0e73f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/4-Object_Detection/MTCNN/main.py b/4-Object_Detection/MTCNN/main.py
@@ -33,7 +33,7 @@ def load_weights(model, weights_file):
     return True
 
 pnet, rnet, onet = PNet(), RNet(), ONet()
-pnet(tf.ones(shape=[1, 224, 224, 3]))
+pnet(tf.ones(shape=[1,  12,  12, 3]))
 rnet(tf.ones(shape=[1,  24,  24 ,3]))
 onet(tf.ones(shape=[1,  48,  48, 3]))
 load_weights(pnet, "./det1.npy"), load_weights(rnet, "./det2.npy"), load_weights(onet, "./det3.npy")

diff --git a/4-Object_Detection/MTCNN/mtcnn.py b/4-Object_Detection/MTCNN/mtcnn.py
@@ -56,13 +56,15 @@ def __init__(self):
         self.dense5_1 = tf.keras.layers.Dense(2, name="conv5-1")
         self.dense5_2 = tf.keras.layers.Dense(4, name="conv5-2")
 
+        self.flatten = tf.keras.layers.Flatten()
+
     def call(self, x, training=False):
         out = self.prelu1(self.conv1(x))
         out = tf.nn.max_pool2d(out, 3, 2, padding="SAME")
         out = self.prelu2(self.conv2(out))
         out = tf.nn.max_pool2d(out, 3, 2, padding="VALID")
         out = self.prelu3(self.conv3(out))
-        out = tf.reshape(out, shape=(out.shape[0], -1))
+        out = self.flatten(out)
         out = self.prelu4(self.dense4(out))
         score = tf.nn.softmax(self.dense5_1(out), -1)
         boxes = self.dense5_2(out)
@@ -91,6 +93,8 @@ def __init__(self):
         self.dense6_2 = tf.keras.layers.Dense(4  , name="conv6-2")
         self.dense6_3 = tf.keras.layers.Dense(10 , name="conv6-3")
 
+        self.flatten = tf.keras.layers.Flatten()
+
     def call(self, x, training=False):
         out = self.prelu1(self.conv1(x))
         out = tf.nn.max_pool2d(out, 3, 2, padding="SAME")
@@ -99,10 +103,13 @@ def call(self, x, training=False):
         out = self.prelu3(self.conv3(out))
         out = tf.nn.max_pool2d(out, 2, 2, padding="SAME")
         out = self.prelu4(self.conv4(out))
-        out = self.dense5(tf.reshape(out, shape=(out.shape[0], -1)))
+
+
+        out = self.dense5(self.flatten(out))
         out = self.prelu5(out)
         score = tf.nn.softmax(self.dense6_1(out))
         boxes = self.dense6_2(out)
         lamks = self.dense6_3(out)
         return boxes, lamks, score
 
+