Merge pull request hunglc007#61 from wooruang/master

Modify some codes to use own train weights for the detect scripts

Author: hunglc007

detect.py

@@ -65,7 +65,11 @@ def main(_argv):
                     bbox_tensor = decode(fm, NUM_CLASS, i)
                     bbox_tensors.append(bbox_tensor)
                 model = tf.keras.Model(input_layer, bbox_tensors)
-                utils.load_weights(model, FLAGS.weights)
+
+                if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
+                    utils.load_weights(model, FLAGS.weights)
+                else:
+                    model.load_weights(FLAGS.weights).expect_partial()
 
         model.summary()
         pred_bbox = model.predict(image_data)

detectvideo.py

@@ -61,7 +61,11 @@ def main(_argv):
                     bbox_tensor = decode(fm, NUM_CLASS, i)
                     bbox_tensors.append(bbox_tensor)
                 model = tf.keras.Model(input_layer, bbox_tensors)
-                utils.load_weights(model, FLAGS.weights)
+                
+                if FLAGS.weights.split(".")[len(FLAGS.weights.split(".")) - 1] == "weights":
+                    utils.load_weights(model, FLAGS.weights)
+                else:
+                    model.load_weights(FLAGS.weights).expect_partial()
 
         model.summary()
     else: