api changes

schemas/api_schema.py

@@ -8,3 +8,5 @@ from pydantic import BaseModel
 class optimization(BaseModel):
     num_class: str
     image_size: str
+    weight_path: str
+    mount_path: str

yolo_optimization.py

@@ -10,9 +10,11 @@ import uvicorn
 app = FastAPI()
 
 class ModelOptimization:
-    def __init__(self, num_class, image_size=416):
+    def __init__(self, num_class, mount_path, weight_path, image_size=416):
         self.num_class = num_class
         self.image_size = image_size
+        self.mount_path = mount_path
+        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}")
@@ -32,16 +34,17 @@ class ModelOptimization:
             logger.info("Successfully changed configurations")
         except Exception as e:
             logger.info(f"Failed to change configurations :  {e}")
+    def transfer_weight_file(self, remote_path, remote_cred):
 
-    def optimize_model(self, weight_path):
+    def optimize_model(self):
         try:
             current_directory = os.getcwd()
             if os.path.exists(os.path.join(current_directory, 'tensorrtx/yolov5/build')):
 
                 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)
+            shutil.copy(self.weight_path, os.path.join(current_directory, 'tensorrtx/yolov5/build'))
+            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}")
@@ -60,6 +63,7 @@ class ModelOptimization:
             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)
 
         except Exception as e:
             logger.info(f"Failed to optimized model :  {e}")
@@ -67,9 +71,9 @@ class ModelOptimization:
 @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 = 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(weight_path = "jk_v5_cam_47.wts")
+    obj.optimize_model()
     return {"message": "successfull"}