api changes

yolo_optimization.py

@@ -36,14 +36,14 @@ class ModelOptimization:
             logger.info(f"Failed to change configurations :  {e}")
     def optimize_model(self):
         try:
-            current_directory = os.getcwd()
-            logger.info(f"Current directory is :  {current_directory}")
-            if os.path.exists(os.path.join(current_directory, 'tensorrtx/yolov5/build')):
-                shutil.rmtree(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            code_directory = "/opt/yolo_model_optimization"
+            logger.info(f"Current directory is :  {code_directory}")
+            if os.path.exists(os.path.join(code_directory, 'tensorrtx/yolov5/build')):
+                shutil.rmtree(os.path.join(code_directory, 'tensorrtx/yolov5/build'))
                 logger.info("Removed existing build directory")
-            os.mkdir(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            os.mkdir(os.path.join(code_directory, 'tensorrtx/yolov5/build'))
             logger.info("Built new build directory")
-            shutil.copy(self.weight_path, os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            shutil.copy(self.weight_path, os.path.join(code_directory, 'tensorrtx/yolov5/build'))
             logger.info("Copied wts file to build directory")
             weight_name_with_extension = os.path.basename(self.weight_path)
             os.chdir('tensorrtx/yolov5/build')
@@ -54,9 +54,9 @@ class ModelOptimization:
             subprocess.run(['make'])
             logger.info("Optimizing model")
             engine_name = "best.engine"
-            subprocess.run(["sudo", "./yolov5", "-s",os.path.join(current_directory, 'tensorrtx/yolov5/build', weight_name_with_extension) , engine_name, "c", "0.33", "0.50"])
-            shutil.move(os.path.join(current_directory, 'tensorrtx/yolov5/build', engine_name), self.mount_path)
-            shutil.move(os.path.join(current_directory, 'tensorrtx/yolov5/build', "libmyplugins.so"), self.mount_path)
+            subprocess.run(["sudo", "./yolov5", "-s",os.path.join(code_directory, 'tensorrtx/yolov5/build', weight_name_with_extension) , engine_name, "c", "0.33", "0.50"])
+            shutil.move(os.path.join(code_directory, 'tensorrtx/yolov5/build', engine_name), self.mount_path)
+            shutil.move(os.path.join(code_directory, 'tensorrtx/yolov5/build', "libmyplugins.so"), self.mount_path)
             logger.info(f"Moved engine file and libmyplugins.so file to {self.mount_path}")
 
         except Exception as e: