SecondPart - master - exenty/UniProject_KorobkoSV

SecondPart
133 строки5 KB
exenty
create: SecondPart
01 июн 2026, 14:43
01 июн 2026, 14:4387aeadb
import matplotlib
import matplotlib.pyplot as plt
import os
import glob
from collections import Counter

# Графики для анализа дорожного покрытия
# Настройка отображения кириллицы
plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['axes.unicode_minus'] = False

# Словарь для сопоставления ID класса и названия
CLASS_NAMES = {
    0: 'D00: Продольная трещина',
    1: 'D10: Поперечная трещина',
    2: 'D20: Сетка трещин',
    3: 'D40: Выбоина',
    4: 'Other: Другой тип'
}

# Цвета для столбцов
COLORS = ['#2E86AB', '#A23B72', '#F18F01', '#C73E1D', '#6A994E']


def count_classes_from_yolo(annotations_dir):

    class_counter = Counter()

    # Поиск всех .txt файлов в директории
    txt_files = glob.glob(os.path.join(annotations_dir, '*.txt'))

    print(f"Найдено файлов аннотаций: {len(txt_files)}")

    for file_path in txt_files:
        try:
            with open(file_path, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    line = line.strip()
                    if not line:
                        continue

                    parts = line.split()
                    if len(parts) < 5:
                        print(
                            f"Предупреждение: некорректная строка {line_num} в файле {os.path.basename(file_path)}")
                        continue

                    try:
                        class_id = int(parts[0])
                        if class_id in CLASS_NAMES:
                            class_counter[class_id] += 1
                        else:
                            print(f"Неизвестный класс {class_id} в файле {os.path.basename(file_path)}")
                    except ValueError:
                        print(f"Ошибка парсинга class_id в строке {line_num} файла {os.path.basename(file_path)}")

        except Exception as e:
            print(f"Ошибка чтения файла {file_path}: {e}")

    return class_counter


def plot_class_distribution(class_counter, save_path=None):

    # Подготовка данных для визуализации
    classes = [CLASS_NAMES[i] for i in sorted(CLASS_NAMES.keys())]
    counts = [class_counter.get(i, 0) for i in sorted(CLASS_NAMES.keys())]
    total = sum(counts)
    percentages = [round(c / total * 100, 1) if total > 0 else 0 for c in counts]

    # Создание фигуры
    fig, ax = plt.subplots(figsize=(12, 7))

    # Построение столбчатой диаграммы
    bars = ax.bar(classes, counts, color=COLORS, edgecolor='black', linewidth=1.2)

    # Добавление подписей значений и процентов
    for bar, count, pct in zip(bars, counts, percentages):
        height = bar.get_height()
        if height > 0:  # Подписывать только ненулевые значения
            ax.text(bar.get_x() + bar.get_width() / 2., height,
                    f'{count:,}\n({pct}%)',
                    ha='center', va='bottom', fontsize=10, fontweight='bold')

    # Оформление осей и заголовка
    ax.set_xlabel('Класс повреждения', fontsize=12, fontweight='bold')
    ax.set_ylabel('Количество объектов', fontsize=12, fontweight='bold')
    ax.set_title('Рисунок 3.1. Распределение объектов повреждений по классам в RDD-2022',
                 fontsize=14, fontweight='bold', pad=20)

    # Настройка сетки
    ax.grid(axis='y', linestyle='--', alpha=0.3)
    ax.set_axisbelow(True)
    plt.xticks(rotation=45, ha='right', fontsize=10)
    plt.yticks(fontsize=10)

    # Добавление общей статистики
    ax.text(0.98, 0.95, f'Всего объектов: {total:,}',
            transform=ax.transAxes, fontsize=11,
            bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.3),
            ha='right', va='top')

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"✓ Диаграмма сохранена: {save_path}")
    plt.show()

    return fig, ax



if __name__ == '__main__':
    ANNOTATIONS_DIR = 'C:/Users/0/Desktop/Курсовая/archivePic/RDD_SPLIT/val/labels'

    # 1. Подсчёт классов
    class_counter = count_classes_from_yolo(ANNOTATIONS_DIR)

    # 2. Вывод краткой статистики в консоль
    print("\nРаспределение классов:")
    print("-" * 50)
    total = sum(class_counter.values())
    for class_id in sorted(CLASS_NAMES.keys()):
        count = class_counter.get(class_id, 0)
        pct = round(count / total * 100, 1) if total > 0 else 0
        print(f"{CLASS_NAMES[class_id]:30s} {count:6,} ({pct}%)")
    print("-" * 50)
    print(f"{'ВСЕГО':30s} {total:6,} (100%)")

    # 3. Построение диаграммы
    print("\nПостроение диаграммы")
    plot_class_distribution(class_counter, save_path='figure_3_1_distribution.png')