task48_images.py - master - exender/coursework-data-analysis

task48_images.py
266 строк12 KB
exender
Create: task49_nlp.py, task48_images.py, task64_timeseries.py, task58_insurance.py
08 июн 2026, 14:44
08 июн 2026, 14:440500eb6
"""
Задача 48: Идентификация средств индивидуальной защиты (СИЗ)
Датасет: PPE Detection Dataset (Roboflow / Kaggle)
Источник: https://universe.roboflow.com/trial-owmex/ppe-detection-qnfpk
         https://www.kaggle.com/datasets/snehilsanyal/construction-site-safety-image-dataset-roboflow

Установка зависимостей:
    pip install pandas numpy matplotlib seaborn pillow opencv-python

Структура папки с датасетом (YOLO format):
    ppe-dataset/
        images/
            train/ val/ test/   ← jpg/png файлы
        labels/
            train/ val/ test/   ← txt файлы (class cx cy w h)

Запуск:
    python task48_images.py
"""

import numpy as np
import pandas as pd
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import seaborn as sns
import os
import json
from pathlib import Path
import warnings
warnings.filterwarnings('ignore')

# Классы PPE датасета
CLASS_NAMES = {
    0: 'Hardhat',
    1: 'Mask',
    2: 'NO-Hardhat',
    3: 'NO-Mask',
    4: 'NO-Safety Vest',
    5: 'Person',
    6: 'Safety Cone',
    7: 'Safety Vest',
    8: 'machinery',
    9: 'vehicle'
}

PPE_COLORS = {
    'Hardhat': '#2ECC71', 'Mask': '#3498DB', 'NO-Hardhat': '#E74C3C',
    'NO-Mask': '#E67E22', 'NO-Safety Vest': '#C0392B', 'Person': '#9B59B6',
    'Safety Cone': '#F39C12', 'Safety Vest': '#1ABC9C',
    'machinery': '#7F8C8D', 'vehicle': '#2C3E50'
}


def load_yolo_labels(labels_dir):
    """Загружает аннотации YOLO из папки labels_dir и возвращает DataFrame."""
    records = []
    labels_path = Path(labels_dir)
    for split in ['train', 'val', 'test']:
        split_dir = labels_path / split
        if not split_dir.exists():
            continue
        for lf in split_dir.glob('*.txt'):
            for line in lf.read_text().strip().split('\n'):
                parts = line.strip().split()
                if len(parts) == 5:
                    cls, cx, cy, w, h = int(parts[0]), *map(float, parts[1:])
                    records.append({'split': split, 'class_id': cls,
                                    'class_name': CLASS_NAMES.get(cls, f'class_{cls}'),
                                    'cx': cx, 'cy': cy, 'w': w, 'h': h,
                                    'area': w * h})
    return pd.DataFrame(records)


def generate_synthetic_annotations():
    """Генерирует синтетические аннотации, отражающие статистику PPE датасета."""
    print("Генерируем синтетические аннотации (датасет не найден)...")
    np.random.seed(42)
    n_images = {'train': 6400, 'val': 800, 'test': 800}
    
    # Распределение объектов по классам (реалистичное)
    class_weights = [0.24, 0.08, 0.12, 0.06, 0.07, 0.31, 0.03, 0.06, 0.02, 0.01]
    records = []
    
    for split, n_img in n_images.items():
        # В среднем 10 объектов на изображение
        n_objects = int(n_img * 10)
        classes = np.random.choice(list(CLASS_NAMES.keys()), size=n_objects, p=class_weights)
        for cls in classes:
            cx = np.random.beta(2, 2)
            cy = np.random.beta(2, 2)
            w = np.random.beta(1.5, 8) * 0.4 + 0.05
            h = np.random.beta(1.5, 6) * 0.5 + 0.05
            records.append({'split': split, 'class_id': cls,
                            'class_name': CLASS_NAMES[cls],
                            'cx': cx, 'cy': cy, 'w': w, 'h': h,
                            'area': w * h})
    return pd.DataFrame(records)


# Загрузка данных
if os.path.isdir('ppe-dataset/labels'):
    df = load_yolo_labels('ppe-dataset/labels')
    print(f"Загружено {len(df)} аннотаций из {df['split'].nunique()} сплитов")
else:
    df = generate_synthetic_annotations()

print(f"\nВсего объектов: {len(df):,}")
print(f"Сплиты: {df['split'].value_counts().to_dict()}")

# ─── 1. Распределение классов ────────────────────────────────────────────────
def plot_class_distribution():
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
    fig.suptitle('Распределение классов датасета PPE Detection', fontsize=14, fontweight='bold')

    # Столбчатая диаграмма
    counts = df['class_name'].value_counts()
    colors = [PPE_COLORS.get(c, 'steelblue') for c in counts.index]
    bars = axes[0].bar(counts.index, counts.values, color=colors, edgecolor='white')
    axes[0].set_title('Количество объектов по классам', fontweight='bold')
    axes[0].set_xlabel('Класс')
    axes[0].set_ylabel('Количество объектов')
    axes[0].tick_params(axis='x', rotation=40)
    for bar, val in zip(bars, counts.values):
        axes[0].text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 100,
                     f'{val:,}', ha='center', va='bottom', fontsize=8)
    axes[0].grid(alpha=0.3, axis='y')

    # Pie chart
    explode = [0.05 if c in ['NO-Hardhat', 'NO-Mask', 'NO-Safety Vest'] else 0 for c in counts.index]
    axes[1].pie(counts.values, labels=counts.index, colors=colors,
                autopct='%1.1f%%', startangle=90, explode=explode,
                textprops={'fontsize': 9})
    axes[1].set_title('Доля классов (%)\n[выделены классы нарушений]', fontweight='bold')

    plt.tight_layout()
    plt.savefig('fig8_class_distribution.png', dpi=150, bbox_inches='tight')
    plt.close()
    print("Рисунок 8 сохранён: fig8_class_distribution.png")

plot_class_distribution()

# ─── 2. Распределение размеров bounding box ──────────────────────────────────
def plot_bbox_sizes():
    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
    fig.suptitle('Анализ размеров bounding box (нормализованные координаты)',
                 fontsize=13, fontweight='bold')

    # Ширина
    axes[0].hist(df['w'], bins=50, color='steelblue', edgecolor='white', alpha=0.85)
    axes[0].axvline(df['w'].median(), color='red', linestyle='--',
                    label=f'Медиана: {df["w"].median():.3f}')
    axes[0].set_title('Ширина bbox (w)', fontweight='bold')
    axes[0].set_xlabel('Нормализованная ширина')
    axes[0].set_ylabel('Частота')
    axes[0].legend()
    axes[0].grid(alpha=0.3)

    # Высота
    axes[1].hist(df['h'], bins=50, color='coral', edgecolor='white', alpha=0.85)
    axes[1].axvline(df['h'].median(), color='navy', linestyle='--',
                    label=f'Медиана: {df["h"].median():.3f}')
    axes[1].set_title('Высота bbox (h)', fontweight='bold')
    axes[1].set_xlabel('Нормализованная высота')
    axes[1].set_ylabel('Частота')
    axes[1].legend()
    axes[1].grid(alpha=0.3)

    # Scatter w vs h
    sample = df.sample(min(3000, len(df)), random_state=42)
    sc_colors = [PPE_COLORS.get(c, 'grey') for c in sample['class_name']]
    axes[2].scatter(sample['w'], sample['h'], c=sc_colors, alpha=0.4, s=10)
    axes[2].set_title('Ширина vs. Высота bbox', fontweight='bold')
    axes[2].set_xlabel('Ширина (w)')
    axes[2].set_ylabel('Высота (h)')
    axes[2].axhline(0.1, color='red', linestyle=':', alpha=0.5, label='Малые объекты')
    axes[2].axvline(0.1, color='red', linestyle=':', alpha=0.5)
    axes[2].legend(fontsize=9)
    axes[2].grid(alpha=0.3)

    plt.tight_layout()
    plt.savefig('fig9_bbox_sizes.png', dpi=150, bbox_inches='tight')
    plt.close()
    print("Рисунок 9 сохранён: fig9_bbox_sizes.png")

plot_bbox_sizes()

# ─── 3. Анализ баланса по сплитам ───────────────────────────────────────────
def plot_split_analysis():
    fig, axes = plt.subplots(1, 3, figsize=(18, 5))
    fig.suptitle('Распределение классов по сплитам', fontsize=13, fontweight='bold')
    
    for ax, split in zip(axes, ['train', 'val', 'test']):
        subset = df[df['split'] == split]
        counts = subset['class_name'].value_counts()
        colors = [PPE_COLORS.get(c, 'steelblue') for c in counts.index]
        ax.bar(counts.index, counts.values, color=colors, edgecolor='white')
        ax.set_title(f'Сплит: {split} ({len(subset):,} объектов)', fontweight='bold')
        ax.tick_params(axis='x', rotation=45)
        ax.set_ylabel('Количество')
        ax.grid(alpha=0.3, axis='y')
    
    plt.tight_layout()
    plt.savefig('fig10_splits.png', dpi=150, bbox_inches='tight')
    plt.close()
    print("Рисунок 10 сохранён: fig10_splits.png")

plot_split_analysis()

# ─── 4. Статистика аннотаций ────────────────────────────────────────────────
print("\n" + "="*60)
print("СТАТИСТИКА АННОТАЦИЙ")
print("="*60)
stats_df = df.groupby('class_name').agg(
    Количество=('class_id', 'count'),
    Ср_ширина=('w', lambda x: f"{x.mean():.3f}"),
    Ср_высота=('h', lambda x: f"{x.mean():.3f}"),
    Ср_площадь=('area', lambda x: f"{x.mean():.4f}"),
).reset_index()
print(stats_df.to_string(index=False))

# ─── 5. Тепловая карта расположения объектов ────────────────────────────────
def plot_heatmap_positions():
    fig, axes = plt.subplots(1, 2, figsize=(14, 5))
    fig.suptitle('Тепловые карты расположения центров объектов на изображении',
                 fontsize=13, fontweight='bold')

    for ax, cls_name in zip(axes, ['Person', 'Hardhat']):
        subset = df[df['class_name'] == cls_name]
        h_map, xedge, yedge = np.histogram2d(
            subset['cx'], subset['cy'], bins=20, range=[[0, 1], [0, 1]])
        im = ax.imshow(h_map.T, origin='lower', cmap='hot',
                       extent=[0, 1, 0, 1], aspect='auto')
        plt.colorbar(im, ax=ax, label='Количество объектов')
        ax.set_title(f'Класс: {cls_name}', fontweight='bold')
        ax.set_xlabel('X (центр bbox)')
        ax.set_ylabel('Y (центр bbox)')

    plt.tight_layout()
    plt.savefig('fig11_heatmap_positions.png', dpi=150, bbox_inches='tight')
    plt.close()
    print("Рисунок 11 сохранён: fig11_heatmap_positions.png")

plot_heatmap_positions()

# ─── 6. Итоговый вывод ──────────────────────────────────────────────────────
print("\n" + "="*60)
print("ИТОГОВЫЕ ВЫВОДЫ")
print("="*60)
total = len(df)
small_bbox = (df['area'] < 0.01).sum()
print(f"Всего объектов: {total:,}")
print(f"Малых объектов (area < 0.01): {small_bbox} ({small_bbox/total*100:.1f}%)")
print(f"  → Необходима мультимасштабная детекция (FPN в YOLOv8)")

violation_classes = ['NO-Hardhat', 'NO-Mask', 'NO-Safety Vest']
viol_count = df[df['class_name'].isin(violation_classes)].shape[0]
print(f"Классы нарушений: {viol_count} ({viol_count/total*100:.1f}%)")
print(f"  → Дисбаланс относительно Person: "
      f"{df[df['class_name']=='Person'].shape[0] / max(viol_count, 1):.1f}:1")
print(f"  → Решение: focal loss α={0.25}, γ=2.0")

print("\n✓ Анализ изображений завершён.")
print("  Файлы: fig8_class_distribution.png, fig9_bbox_sizes.png,")
print("         fig10_splits.png, fig11_heatmap_positions.png")