From 9fc471380de58a9b71ab126f90cd63d512102f22 Mon Sep 17 00:00:00 2001
From: aimspot <som_nncs@list.ru>
Date: Wed, 14 Feb 2024 13:04:06 +0200
Subject: [PATCH] fix

---
 .../train_model/scripts/yolov5_train.py       | 21 +++++++++----------
 ODRS/utils/dataset_info.py                    |  7 +++++--
 ODRS/utils/ml_plot.py                         |  3 ++-
 3 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/ODRS/train_utils/train_model/scripts/yolov5_train.py b/ODRS/train_utils/train_model/scripts/yolov5_train.py
index 1c78fd2..2a27d85 100755
--- a/ODRS/train_utils/train_model/scripts/yolov5_train.py
+++ b/ODRS/train_utils/train_model/scripts/yolov5_train.py
@@ -14,12 +14,12 @@
 #     opt.name = 'exp'
 #     train.main(opt)
 
-# ----------------------------------------------OLD version
+
 def train_V5(IMG_SIZE, BATCH_SIZE, EPOCHS, CONFIG_PATH, MODEL_PATH, GPU_COUNT, SELECT_GPU):
     """
     Runs yolov5 training using the parameters specified in the config.
 
-    
+
     :param IMG_SIZE: Size of input images as integer or w,h.
     :param BATCH_SIZE: Batch size for training.
     :param EPOCHS: Number of epochs to train for.
@@ -33,15 +33,14 @@ def train_V5(IMG_SIZE, BATCH_SIZE, EPOCHS, CONFIG_PATH, MODEL_PATH, GPU_COUNT, S
     train_script_path = str(Path(file.parents[1]) / 'models' / 'yolov5' / 'train.py')
 
     full_command = (
-        command +
-        f" {train_script_path}" +
-        f" --img {IMG_SIZE}" +
-        f" --batch {BATCH_SIZE}" +
-        f" --epochs {EPOCHS}" +
-        f" --data {CONFIG_PATH}" +
-        f" --cfg {MODEL_PATH}" +
-        f" --device {SELECT_GPU}" +
-        f" --project {CONFIG_PATH.parent}" +
+        f"{command} {train_script_path}"
+        f" --img {IMG_SIZE}"
+        f" --batch {BATCH_SIZE}"
+        f" --epochs {EPOCHS}"
+        f" --data {CONFIG_PATH}"
+        f" --cfg {MODEL_PATH}"
+        f" --device {SELECT_GPU}"
+        f" --project {CONFIG_PATH.parent}"
         f" --name exp"
     )
     os.system(full_command)
diff --git a/ODRS/utils/dataset_info.py b/ODRS/utils/dataset_info.py
index 6b3da08..38f33f8 100644
--- a/ODRS/utils/dataset_info.py
+++ b/ODRS/utils/dataset_info.py
@@ -1,15 +1,18 @@
-import cv2
 import os
-import numpy as np
 import sys
 from loguru import logger
 from pathlib import Path
+import cv2
+import numpy as np
+
 project_dir = os.path.dirname(os.path.abspath(__file__))
 sys.path.append(os.path.dirname(os.path.dirname(project_dir)))
+
 from ODRS.utils.ml_plot import plot_class_balance
 from ODRS.utils.ml_utils import dumpCSV
 
 
+
 def load_class_names(classes_file):
     """ Загрузка названий классов из файла. """
     with open(classes_file, 'r') as file:
diff --git a/ODRS/utils/ml_plot.py b/ODRS/utils/ml_plot.py
index b6dc211..1222280 100644
--- a/ODRS/utils/ml_plot.py
+++ b/ODRS/utils/ml_plot.py
@@ -3,6 +3,7 @@
 import random
 from pathlib import Path
 
+
 def plot_class_balance(labels, output_path):
     """ Построение и сохранение графика баланса классов с наклоненными метками и вывод среднего значения. """
     class_counts = Counter(labels)
@@ -15,4 +16,4 @@ def plot_class_balance(labels, output_path):
     plt.title('Class balance')
     plt.xticks(rotation=90)
     plt.tight_layout()
-    plt.savefig(output_file)
\ No newline at end of file
+    plt.savefig(output_file)