added file copy logic

yolo_optimization.py

@@ -5,9 +5,10 @@ import shutil
 
 
 class ModelOptimization:
-    def __init__(self, num_class, image_size=416):
+    def __init__(self, num_class, weight_path, 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}")
@@ -28,10 +29,10 @@ class ModelOptimization:
         except Exception as e:
             logger.info(f"Failed to change configurations :  {e}")
 
-    def optimize_model(self, weight_path):
+    def optimize_model(self):
         try:
-            shutil.copy(weight_path, 'yolov5/build')
-            weight_name_with_extension = os.path.basename(weight_path)
+            shutil.copy(self.weight_path, 'yolov5/build')
+            weight_name_with_extension = os.path.basename(self.weight_path)
             weight_name, extension = os.path.splitext(weight_name_with_extension)
             current_directory = os.getcwd()
             logger.info(f"Current directory is :  {current_directory}")
@@ -51,6 +52,8 @@ class ModelOptimization:
             logger.info(f"Failed to optimized model :  {e}")
 
 
-# obj = ModelOptimization(num_class=2, image_size=244)
-# obj.change_configurations()
+obj = ModelOptimization(num_class=1,weight_path="/home/ilens/cam_42_best.wts", image_size=416)
+obj.change_configurations()
+obj.optimize_model()
+