api changes

yolo_optimization.py

@@ -37,23 +37,17 @@ class ModelOptimization:
     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'))
+                logger.info("Removed existing build directory")
             os.mkdir(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            logger.info("Built new build directory")
             shutil.copy(self.weight_path, os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            logger.info("Copied wts file to build directory")
             weight_name_with_extension = os.path.basename(self.weight_path)
-            weight_name, extension = os.path.splitext(weight_name_with_extension)
-
-            logger.info(f"Current directory is :  {current_directory}")
-            # build_files = glob.glob("tensorrtx/yolov5/build/*")
-            # for file in build_files:
-            #     os.remove(file)
-
-            # build_path = os.path.join(current_directory, "yolov5", "build")
-            # os.mkdir(build_path)
-            # logger.info(f"Created build folder")
             os.chdir('tensorrtx/yolov5/build')
+            logger.info("Changed to build directory")
             logger.info("Running CMake command")
             subprocess.run(['cmake', '..'])
             logger.info("Running Make command")
@@ -62,13 +56,14 @@ class ModelOptimization:
             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)
+            logger.info(f"Moved engine file and libmyplugins.so file to {self.mount_path}")
 
         except Exception as e:
             logger.info(f"Failed to optimized model :  {e}")
 
 @app.post("/optimize")
 async def root(content: optimization):
-    # print(content.dict())
     obj = ModelOptimization(num_class=int(content.num_class), image_size=int(content.image_size), weight_path = content.weight_path, mount_path=content.mount_path)
     obj.change_configurations()
     obj.optimize_model()