Print summary when --profile option is set

axinc-ai · Aug 1, 2024 · 6c0ffa1 · 6c0ffa1
1 parent 53b2a7b
commit 6c0ffa1
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/road_detection/hybridnets/hybridnets.py b/road_detection/hybridnets/hybridnets.py
@@ -126,6 +126,8 @@ def recognize_from_image():
     t0 = time.time()
 
     model = ailia.Net(None,args.model_name, env_id=args.env_id)
+    if args.profile:
+        model.set_profile_mode(True)
 
     color_list = standard_to_bgr(STANDARD_COLORS)
 
@@ -168,6 +170,9 @@ def recognize_from_image():
             cap.release()
             out_stream.release()
 
+    if args.profile:
+        print(model.get_summary())
+
     logger.info('Script finished successfully.')
 
 def recognize_from_video():