Commit 70490026 authored by Sikhin VC's avatar Sikhin VC

added file copy logic

parent 0b4e78cf
import os
import subprocess
from loguru import logger
import shutil
class ModelOptimization:
......@@ -27,9 +28,11 @@ class ModelOptimization:
except Exception as e:
logger.info(f"Failed to change configurations : {e}")
def optimize_model(self, ):
def optimize_model(self, weight_path):
try:
shutil.copy(weight_path, 'yolov5/build')
weight_name_with_extension = os.path.basename(weight_path)
weight_name, extension = os.path.splitext(weight_name_with_extension)
current_directory = os.getcwd()
logger.info(f"Current directory is : {current_directory}")
build_path = os.path.join(current_directory, "yolov5", "build")
......@@ -41,7 +44,8 @@ class ModelOptimization:
logger.info("Running Make command")
subprocess.run(['make'])
logger.info("Optimizing model")
subprocess.run(["sudo", "./yolov5", "-s", "jk_v5_cam_47.wts", "jk_v5_cam_47.engine", "c", "0.33", "0.50"])
engine_name = weight_name + ".engine"
subprocess.run(["sudo", "./yolov5", "-s", weight_name_with_extension, engine_name, "c", "0.33", "0.50"])
except Exception as e:
logger.info(f"Failed to optimized model : {e}")
......
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