engine optimization service

README.md

-# yolo_model_optimization
-

schemas/api_schema.py

+from __future__ import annotations
+
+from typing import Any, List, Optional
+
+from pydantic import BaseModel
+
+
+class optimization(BaseModel):
+    num_class: str
+    image_size: str

yolov5/yololayer.h → tensorrtx/yolov5/yololayer.h

@@ -17,9 +17,9 @@ namespace Yolo
         float anchors[CHECK_COUNT * 2];
     };
     static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000;
-    static constexpr int CLASS_NUM = 80;
-    static constexpr int INPUT_H = 640;  // yolov5's input height and width must be divisible by 32.
-    static constexpr int INPUT_W = 640;
+    static constexpr int CLASS_NUM = 2;
+    static constexpr int INPUT_H = 416;
+    static constexpr int INPUT_H = 416;
 
     static constexpr int LOCATIONS = 4;
     struct alignas(float) Detection {

yolo_optimization.py

+import glob
 import os
 import subprocess
 from loguru import logger
 import shutil
+from fastapi import FastAPI
+from schemas.api_schema import optimization
 
+app = FastAPI()
 
 class ModelOptimization:
-    def __init__(self, num_class, weight_path, image_size=416):
+    def __init__(self, num_class, image_size=416):
         self.num_class = num_class
         self.image_size = image_size
-        self.weight_path = weight_path
 
     def change_configurations(self):
         logger.info(f"Provided number of classes and image size are :  {self.num_class} and {self.image_size}")
         try:
 
-            with open('yolov5/yololayer.h', 'r') as file:
+            with open('tensorrtx/yolov5/yololayer.h', 'r') as file:
                 # read a list of lines into data
                 data = file.readlines()
 
             data[19] = f"    static constexpr int CLASS_NUM = {self.num_class};\n"
             data[20] = f"    static constexpr int INPUT_H = {self.image_size};\n"
-            data[21] = f"    static constexpr int INPUT_W = {self.image_size};\n"
+            data[21] = f"    static constexpr int INPUT_H = {self.image_size};\n"
 
             # and write everything back
-            with open('yolov5/yololayer.h', 'w') as file:
+            with open('tensorrtx/yolov5/yololayer.h', 'w') as file:
                 file.writelines(data)
             logger.info("Successfully changed configurations")
         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, 'tensorrtx/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")
-            if os.path.isdir('yolov5/build'):
-                logger.info("build directory exists. Removing build directory!!")
-                shutil.rmtree('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 = f"{current_directory}/yolov5/build/"
-            shutil.copy(self.weight_path, src)
-            logger.info(f"Created build folder")
-            os.chdir('yolov5/build')
+            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("Running CMake command")
             subprocess.run(['cmake', '..'])
             logger.info("Running Make command")
             subprocess.run(['make'])
             logger.info("Optimizing model")
-            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"])
 
         except Exception as e:
             logger.info(f"Failed to optimized model :  {e}")
 
-
-obj = ModelOptimization(num_class=1,weight_path="/home/ilens/cam_42_best.wts", image_size=416)
-obj.change_configurations()
-obj.optimize_model()
+@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))
+    obj.change_configurations()
+    obj.optimize_model(weight_path = "jk_v5_cam_47.wts")
+    return {"message": "successfull"}