abc

yolo_optimization.py

@@ -36,7 +36,9 @@ class ModelOptimization:
     def optimize_model(self, weight_path):
         try:
             current_directory = os.getcwd()
-            shutil.rmtree(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            if os.path.exists(os.path.join(current_directory, 'tensorrtx/yolov5/build')):
+
+                shutil.rmtree(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
             os.mkdir(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
             shutil.copy(weight_path, os.path.join(current_directory, 'tensorrtx/yolov5/build'))
             weight_name_with_extension = os.path.basename(weight_path)