Skip to content

Y
yolo_model_optimization
Commit a9f78581 authored by Sikhin VC's avatar Sikhin VC
Browse files
Options

abc

parent e603f388
Hide whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
yolo_optimization.py
View file @ a9f78581
...@@ -36,7 +36,9 @@ class ModelOptimization: ...@@ -36,7 +36,9 @@ class ModelOptimization:
def optimize_model(self, weight_path): def optimize_model(self, weight_path):
try: try:
current_directory = os.getcwd() 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')) os.mkdir(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
shutil.copy(weight_path, 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) weight_name_with_extension = os.path.basename(weight_path)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment