abc

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 = 2;
-    static constexpr int INPUT_H = 416;
+    static constexpr int CLASS_NUM = 1;
     static constexpr int INPUT_H = 416;
+    static constexpr int INPUT_W = 416;
 
     static constexpr int LOCATIONS = 4;
     struct alignas(float) Detection {

yolo_optimization.py

@@ -5,6 +5,7 @@ from loguru import logger
 import shutil
 from fastapi import FastAPI
 from schemas.api_schema import optimization
+import uvicorn
 
 app = FastAPI()
 
@@ -23,7 +24,7 @@ class ModelOptimization:
 
             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_H = {self.image_size};\n"
+            data[21] = f"    static constexpr int INPUT_W = {self.image_size};\n"
 
             # and write everything back
             with open('tensorrtx/yolov5/yololayer.h', 'w') as file:
@@ -34,14 +35,17 @@ class ModelOptimization:
 
     def optimize_model(self, weight_path):
         try:
-            shutil.copy(weight_path, 'tensorrtx/yolov5/build')
+            current_directory = os.getcwd()
+            shutil.rmtree(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            os.mkdir(os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            shutil.copy(weight_path, os.path.join(current_directory, '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_files = glob.glob("tensorrtx/yolov5/build/*")
-            for file in build_files:
-                os.remove(file)
+            # 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)
@@ -53,7 +57,7 @@ class ModelOptimization:
             subprocess.run(['make'])
             logger.info("Optimizing model")
             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",os.path.join(current_directory, 'tensorrtx/yolov5/build', weight_name_with_extension) , engine_name, "c", "0.33", "0.50"])
 
         except Exception as e:
             logger.info(f"Failed to optimized model :  {e}")
@@ -67,3 +71,7 @@ async def root(content: optimization):
     return {"message": "successfull"}
 
 
+if __name__ == '__main__':
+    uvicorn.run(app, port=8080, host='localhost')
+
+