Fix plot

ODRS/train_utils/config/custom_config.yaml

@@ -4,7 +4,7 @@ DATA_PATH: Warp-D
 EPOCHS: 2
 GPU_COUNT: 1
 IMG_SIZE: 200
-MODEL: faster-rcnn
+MODEL: yolov5l
 SELECT_GPU: 0
 
 CONFIG_PATH: dataset.yaml

ODRS/utils/dataset_info.py

@@ -3,8 +3,7 @@
 import numpy as np
 from loguru import logger
 from pathlib import Path
-from collections import Counter
-import matplotlib.pyplot as plt
+from ODRS.utils.ml_plot import plot_class_balance
 
 def load_class_names(classes_file):
     """ Загрузка названий классов из файла. """
@@ -31,20 +30,6 @@ def load_yolo_labels(data_path, class_names):
     return labels
 
 
-def plot_class_balance(labels, output_path):
-    """ Построение и сохранение графика баланса классов с наклоненными метками и вывод среднего значения. """
-    class_counts = Counter(labels)
-    output_file = output_path / 'Class_balance.png'
-
-    plt.bar(class_counts.keys(), class_counts.values())
-    plt.xlabel('Classes')
-    plt.ylabel('Number of instances')
-    plt.title('Class balance')
-    plt.xticks(rotation=90)
-    plt.tight_layout()
-    plt.savefig(output_file)
-    # plt.show()
-
 def find_images(data_path):
     supported_extensions = {".jpg", ".jpeg", ".png"}
     image_paths = []

ODRS/utils/ml_plot.py

@@ -0,0 +1,17 @@
+from collections import Counter
+import matplotlib.pyplot as plt
+import random
+
+def plot_class_balance(labels, output_path):
+    """ Построение и сохранение графика баланса классов с наклоненными метками и вывод среднего значения. """
+    class_counts = Counter(labels)
+    output_file = output_path / 'Class_balance.png'
+    colors = [f'#{random.randint(0, 0xFFFFFF):06x}' for _ in class_counts.keys()]
+
+    plt.bar(class_counts.keys(), class_counts.values(), color=colors)
+    plt.xlabel('Classes')
+    plt.ylabel('Number of instances')
+    plt.title('Class balance')
+    plt.xticks(rotation=90)
+    plt.tight_layout()
+    plt.savefig(output_file)