Chapter_3.py - master - vasilenko/Chapter_1

Chapter_3.py
130 строк6 KB
vasilenko
Create: Chapter_3.py
11 июн 2026, 15:59
11 июн 2026, 15:59f6308f1
# -*- coding: utf-8 -*-
"""
Раздел 3 курсового проекта: первичный анализ набора изображений.
Набор данных: Construction Site Safety Image Dataset (Kaggle, формат YOLOv8).
Скрипт считает объём подвыборок, баланс классов, пустые аннотации,
проверяет размеры изображений и показывает примеры.
"""

from pathlib import Path
from collections import Counter

import yaml                     # чтение файла data.yaml
from PIL import Image           # работа с изображениями
import matplotlib.pyplot as plt

# Путь к папке датасета: внутри должны быть train/, valid/, test/ и data.yaml
DATA_PATH = Path("images_dataset_work/css-data")

SPLITS = ["train", "valid", "test"]

# ============================================
# ЧАСТЬ 1 — Загрузка датасета изображений
# ============================================
# В файле data.yaml перечислены классы и пути к подвыборкам.
with open(DATA_PATH / "data.yaml", encoding="utf-8") as f:
    meta = yaml.safe_load(f)

class_names = meta["names"]
print("=== Метаданные (data.yaml) ===")
print("Число классов:", meta["nc"])
print("Классы:", class_names, "\n")

# Считаем, сколько изображений в каждой подвыборке
counts = {}
for split in SPLITS:
    imgs = list((DATA_PATH / split / "images").glob("*.jpg"))
    counts[split] = len(imgs)
total = sum(counts.values())

print("=== Разделение датасета на подвыборки (таблица 6) ===")
for split in SPLITS:
    print(f"  {split:6s}: {counts[split]:5d} изображений "
          f"({counts[split] / total * 100:.1f} %)")
print(f"  всего : {total} изображений\n")

# ============================================
# ЧАСТЬ 2 — Предобработка / проверка изображений
# ============================================
# Обучения в этой главе нет, поэтому проверяем только формат:
# все ли изображения имеют размер 640x640 пикселей.
sizes = Counter()
for split in SPLITS:
    for img_path in (DATA_PATH / split / "images").glob("*.jpg"):
        with Image.open(img_path) as im:
            sizes[im.size] += 1
print("=== Размеры изображений ===")
for size, n in sizes.most_common():
    print(f"  {size[0]}x{size[1]}: {n} шт.")
print()

# ============================================
# ЧАСТЬ 3 — Основная логика: анализ разметки
# ============================================
# Формат YOLOv8: для каждой картинки есть текстовый файл в labels/,
# каждая строка = один объект: «номер_класса x y w h».
objects = Counter()       # сколько объектов каждого класса
empty_files = 0           # файлы аннотаций без единого объекта
label_files = 0

for split in SPLITS:
    for lab in (DATA_PATH / split / "labels").glob("*.txt"):
        label_files += 1
        lines = [s for s in lab.read_text().splitlines() if s.strip()]
        if not lines:
            empty_files += 1
            continue
        for line in lines:
            cls_id = int(line.split()[0])   # первый столбец — номер класса
            objects[cls_id] += 1

total_obj = sum(objects.values())
print("=== Распределение объектов по классам (таблица 7) ===")
for cls_id in sorted(objects):
    name = class_names[cls_id]
    n = objects[cls_id]
    print(f"  {cls_id:2d} {name:15s}: {n:6d} ({n / total_obj * 100:5.1f} %)")
print(f"\nВсего размеченных объектов: {total_obj}")
print(f"В среднем объектов на изображение: {total_obj / total:.1f}")
print(f"Пустых файлов аннотаций: {empty_files} "
      f"({empty_files / label_files * 100:.1f} %)\n")

# ============================================
# ЧАСТЬ 4 — Вывод результатов (графики)
# ============================================

#разделение датасета на подвыборки
plt.figure(figsize=(6, 4))
plt.bar(counts.keys(), counts.values(), color=["#4c72b0", "#dd8452", "#55a868"])
for i, (split, n) in enumerate(counts.items()):
    plt.text(i, n, str(n), ha="center", va="bottom")
plt.title("Разделение датасета на подвыборки")
plt.ylabel("Число изображений")
plt.tight_layout()
plt.savefig("ris31_podvyborki.png", dpi=150)

#распределение объектов по классам
names = [class_names[c] for c in sorted(objects)]
vals = [objects[c] for c in sorted(objects)]
plt.figure(figsize=(9, 5))
plt.barh(names, vals, color="#4c72b0")
for y, v in enumerate(vals):
    plt.text(v, y, f" {v}", va="center", fontsize=8)
plt.title("Распределение объектов по классам")
plt.xlabel("Число объектов")
plt.tight_layout()
plt.savefig("ris32_klassy.png", dpi=150)

#примеры типичных изображений (первые 5 из train)
examples = sorted((DATA_PATH / "train" / "images").glob("*.jpg"))[:5]
fig, axes = plt.subplots(1, len(examples), figsize=(15, 4))
for ax, img_path in zip(axes, examples):
    ax.imshow(Image.open(img_path))
    ax.set_title(img_path.name, fontsize=6)
    ax.axis("off")
fig.suptitle("Примеры типичных изображений датасета")
plt.tight_layout()
plt.savefig("ris33_primery.png", dpi=150)

print("Готово! Графики сохранены в файлы ris31...ris33 (.png).")
plt.show()