abc

edge_engine/model_optimizer.py

@@ -36,27 +36,29 @@ class ModelOptimization:
             current_directory = os.getcwd()
 
             logger.info(f"Current directory is :  {current_directory}")
-            build_path = os.path.join("/app/edge_engine", "yolov5", "build")
-            if os.path.isdir('/app/edge_engine/yolov5/build'):
-                logger.info("build directory exists. Removing build directory!!")
-                shutil.rmtree('/app/edge_engine/yolov5/build')
-            os.mkdir(build_path)
-            weight_name_with_extension = os.path.basename(self.weight_path)
-            weight_name, extension = os.path.splitext(weight_name_with_extension)
-            src = "/app/edge_engine/yolov5/build"
-            shutil.copy(self.weight_path, src)
-            logger.info(f"Created build folder")
+            # build_path = os.path.join("/app/edge_engine", "yolov5", "build")
+            # if os.path.isdir('/app/edge_engine/yolov5/build'):
+            #     logger.info("build directory exists. Removing build directory!!")
+            #     shutil.rmtree('/app/edge_engine/yolov5/build')
+            # os.mkdir(build_path)
+            # weight_name_with_extension = os.path.basename(self.weight_path)
+            # weight_name, extension = os.path.splitext(weight_name_with_extension)
+            # src = "/app/edge_engine/yolov5/build"
+            # shutil.copy(self.weight_path, src)
+            # logger.info(f"Created build folder")
             os.chdir('/app/edge_engine/yolov5/build')
             logger.info("Running CMake command")
             subprocess.run(['cmake', '..'])
             logger.info("Running Make command")
             subprocess.run(['make'])
             logger.info("Optimizing model")
-            engine_name = weight_name + ".engine"
-            subprocess.run(["sudo", "./yolov5", "-s", weight_name_with_extension, engine_name, "c", "0.33", "0.50"])
-            engine_file_path = os.path.join(src, engine_name)
-
-            shutil.copy(engine_file_path, "/app/data/")
+            # engine_name = weight_name + ".engine"
+            engine_name = "best.engine"
+            # subprocess.run(["sudo", "./yolov5", "-s", weight_name_with_extension, engine_name, "c", "0.33", "0.50"])
+            subprocess.run(["sudo", "./yolov5", "-s", "best.wts", "best.engine", "c", "0.33", "0.50"])
+            # engine_file_path = os.path.join(src, engine_name)
+            #
+            # shutil.copy(engine_file_path, "/app/data/")
 
         except Exception as e:
             logger.info(f"Failed to optimized model :  {e}")