Add MobileNetV3 models support in inference scripts

qfgaohao · Apr 5, 2020 · 5e5a756 · 5e5a756
1 parent eb1205f
commit 5e5a756
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 7 deletions.
diff --git a/eval_ssd.py b/eval_ssd.py
@@ -13,6 +13,7 @@
 import logging
 import sys
 from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite, create_mobilenetv2_ssd_lite_predictor
+from vision.ssd.mobilenetv3_ssd_lite import create_mobilenetv3_large_ssd_lite, create_mobilenetv3_small_ssd_lite
 
 
 parser = argparse.ArgumentParser(description="SSD Evaluation on VOC Dataset.")
@@ -141,10 +142,14 @@ def compute_average_precision_per_class(num_true_cases, gt_boxes, difficult_case
         net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
     elif args.net == 'mb2-ssd-lite':
         net = create_mobilenetv2_ssd_lite(len(class_names), width_mult=args.mb2_width_mult, is_test=True)
+    elif args.net == 'mb3-large-ssd-lite':
+        net = create_mobilenetv3_large_ssd_lite(len(class_names), is_test=True)
+    elif args.net == 'mb3-small-ssd-lite':
+        net = create_mobilenetv3_small_ssd_lite(len(class_names), is_test=True)
     else:
         logging.fatal("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.")
         parser.print_help(sys.stderr)
-        sys.exit(1)  
+        sys.exit(1)
 
     timer.start("Load Model")
     net.load(args.trained_model)
@@ -158,7 +163,7 @@ def compute_average_precision_per_class(num_true_cases, gt_boxes, difficult_case
         predictor = create_mobilenetv1_ssd_lite_predictor(net, nms_method=args.nms_method, device=DEVICE)
     elif args.net == 'sq-ssd-lite':
         predictor = create_squeezenet_ssd_lite_predictor(net,nms_method=args.nms_method, device=DEVICE)
-    elif args.net == 'mb2-ssd-lite':
+    elif args.net == 'mb2-ssd-lite' or args.net == "mb3-large-ssd-lite" or args.net == "mb3-small-ssd-lite":
         predictor = create_mobilenetv2_ssd_lite_predictor(net, nms_method=args.nms_method, device=DEVICE)
     else:
         logging.fatal("The net type is wrong. It should be one of vgg16-ssd, mb1-ssd and mb1-ssd-lite.")
@@ -212,6 +217,3 @@ def compute_average_precision_per_class(num_true_cases, gt_boxes, difficult_case
         print(f"{class_name}: {ap}")
 
     print(f"\nAverage Precision Across All Classes:{sum(aps)/len(aps)}")
-
-
-
diff --git a/run_ssd_example.py b/run_ssd_example.py
@@ -3,6 +3,7 @@
 from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite, create_mobilenetv1_ssd_lite_predictor
 from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite, create_squeezenet_ssd_lite_predictor
 from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite, create_mobilenetv2_ssd_lite_predictor
+from vision.ssd.mobilenetv3_ssd_lite import create_mobilenetv3_large_ssd_lite, create_mobilenetv3_small_ssd_lite
 from vision.utils.misc import Timer
 import cv2
 import sys
@@ -26,6 +27,10 @@
     net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
 elif net_type == 'mb2-ssd-lite':
     net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
+elif net_type == 'mb3-large-ssd-lite':
+    net = create_mobilenetv3_large_ssd_lite(len(class_names), is_test=True)
+elif net_type == 'mb3-small-ssd-lite':
+    net = create_mobilenetv3_small_ssd_lite(len(class_names), is_test=True)
 elif net_type == 'sq-ssd-lite':
     net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
 else:
@@ -39,7 +44,7 @@
     predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
 elif net_type == 'mb1-ssd-lite':
     predictor = create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200)
-elif net_type == 'mb2-ssd-lite':
+elif net_type == 'mb2-ssd-lite' or net_type == "mb3-large-ssd-lite" or net_type == "mb3-small-ssd-lite":
     predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200)
 elif net_type == 'sq-ssd-lite':
     predictor = create_squeezenet_ssd_lite_predictor(net, candidate_size=200)

diff --git a/run_ssd_live_demo.py b/run_ssd_live_demo.py
@@ -3,6 +3,7 @@
 from vision.ssd.mobilenetv1_ssd_lite import create_mobilenetv1_ssd_lite, create_mobilenetv1_ssd_lite_predictor
 from vision.ssd.squeezenet_ssd_lite import create_squeezenet_ssd_lite, create_squeezenet_ssd_lite_predictor
 from vision.ssd.mobilenet_v2_ssd_lite import create_mobilenetv2_ssd_lite, create_mobilenetv2_ssd_lite_predictor
+from vision.ssd.mobilenetv3_ssd_lite import create_mobilenetv3_large_ssd_lite, create_mobilenetv3_small_ssd_lite
 from vision.utils.misc import Timer
 import cv2
 import sys
@@ -33,6 +34,10 @@
     net = create_mobilenetv1_ssd_lite(len(class_names), is_test=True)
 elif net_type == 'mb2-ssd-lite':
     net = create_mobilenetv2_ssd_lite(len(class_names), is_test=True)
+elif net_type == 'mb3-large-ssd-lite':
+    net = create_mobilenetv3_large_ssd_lite(len(class_names), is_test=True)
+elif net_type == 'mb3-small-ssd-lite':
+    net = create_mobilenetv3_small_ssd_lite(len(class_names), is_test=True)
 elif net_type == 'sq-ssd-lite':
     net = create_squeezenet_ssd_lite(len(class_names), is_test=True)
 else:
@@ -46,7 +51,7 @@
     predictor = create_mobilenetv1_ssd_predictor(net, candidate_size=200)
 elif net_type == 'mb1-ssd-lite':
     predictor = create_mobilenetv1_ssd_lite_predictor(net, candidate_size=200)
-elif net_type == 'mb2-ssd-lite':
+elif net_type == 'mb2-ssd-lite' or net_type == "mb3-large-ssd-lite" or net_type == "mb3-small-ssd-lite":
     predictor = create_mobilenetv2_ssd_lite_predictor(net, candidate_size=200)
 elif net_type == 'sq-ssd-lite':
     predictor = create_squeezenet_ssd_lite_predictor(net, candidate_size=200)