jittor-segment-A/data_analysis.py

import os
import numpy as np
import matplotlib.pyplot as plt
import cv2
from collections import Counter
import random
from PIL import Image
import pandas as pd
import seaborn as sns

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

# 数据集路径
TRAIN_SET_PATH = 'TrainSet'
TEST_SET_PATH = 'TestSetA'

# 标签文件路径
train_label_path = os.path.join(TRAIN_SET_PATH, 'labels', 'train.txt')
val_label_path = os.path.join(TRAIN_SET_PATH, 'labels', 'val.txt')
trainval_label_path = os.path.join(TRAIN_SET_PATH, 'labels', 'trainval.txt')

# 图像路径
train_image_dir = os.path.join(TRAIN_SET_PATH, 'images', 'train')
test_image_dir = TEST_SET_PATH

# 读取标签文件
def read_label_file(file_path):
    labels = {}
    with open(file_path, 'r') as f:
        for line in f:
            img_name, label = line.strip().split()
            labels[img_name] = int(label)
    return labels

# 读取训练集和验证集标签
train_labels = read_label_file(train_label_path)
val_labels = read_label_file(val_label_path)
trainval_labels = read_label_file(trainval_label_path)

# 统计各个类别的数量
def count_classes(labels):
    return Counter(labels.values())

# 获取图像尺寸分布
def get_image_sizes(image_dir, sample_size=100):
    image_files = os.listdir(image_dir)
    if len(image_files) > sample_size:
        image_files = random.sample(image_files, sample_size)
    
    sizes = []
    for img_file in image_files:
        img_path = os.path.join(image_dir, img_file)
        try:
            img = Image.open(img_path)
            sizes.append(img.size)
        except Exception as e:
            print(f"Error processing {img_file}: {e}")
    
    return sizes

# 可视化类别分布
def plot_class_distribution(train_counts, val_counts, title="类别分布"):
    plt.figure(figsize=(12, 6))
    
    # 创建类别标签映射
    class_names = {
        0: "类别0",
        1: "类别1",
        2: "类别2",
        3: "类别3",
        4: "类别4",
        5: "类别5"
    }
    
    # 准备数据
    classes = list(set(list(train_counts.keys()) + list(val_counts.keys())))
    # 计算每个类别的总样本数（训练集+验证集）
    total_counts = {cls: train_counts.get(cls, 0) + val_counts.get(cls, 0) for cls in classes}
    # 按照样本总数从大到小排序类别
    classes = sorted(classes, key=lambda x: total_counts[x], reverse=True)
    train_values = [train_counts.get(cls, 0) for cls in classes]
    val_values = [val_counts.get(cls, 0) for cls in classes]
    
    # 创建DataFrame
    df = pd.DataFrame({
        '训练集': train_values,
        '验证集': val_values
    }, index=[class_names.get(cls, f"类别{cls}") for cls in classes])
    
    # 绘制堆叠柱状图
    ax = df.plot(kind='bar', stacked=False, figsize=(12, 6))
    plt.title(title, fontsize=16)
    plt.xlabel('类别', fontsize=14)
    plt.ylabel('样本数量', fontsize=14)
    plt.xticks(rotation=0)
    
    # 添加数值标签
    for container in ax.containers:
        ax.bar_label(container, fmt='%d')
    
    plt.tight_layout()
    plt.savefig('类别分布.png', dpi=300, bbox_inches='tight')
    plt.show()

# 可视化图像尺寸分布
def plot_image_size_distribution(sizes):
    # 提取宽度和高度
    widths, heights = zip(*sizes)
    
    plt.figure(figsize=(12, 6))
    
    # 创建散点图
    plt.scatter(widths, heights, alpha=0.5)
    plt.title('图像尺寸分布', fontsize=16)
    plt.xlabel('宽度 (像素)', fontsize=14)
    plt.ylabel('高度 (像素)', fontsize=14)
    plt.grid(True, linestyle='--', alpha=0.7)
    
    # 添加常见分辨率标记
    common_resolutions = [(640, 480), (800, 600), (1024, 768), (1280, 720), (1920, 1080)]
    for w, h in common_resolutions:
        plt.plot(w, h, 'ro', markersize=10, alpha=0.3)
        plt.annotate(f'{w}x{h}', (w, h), xytext=(10, 10), textcoords='offset points')
    
    plt.tight_layout()
    plt.savefig('图像尺寸分布.png', dpi=300, bbox_inches='tight')
    plt.show()

# 显示每个类别的样本图像
def show_sample_images(image_dir, labels, num_per_class=3):
    # 按类别分组图像
    class_images = {}
    for img_name, label in labels.items():
        if label not in class_images:
            class_images[label] = []
        class_images[label].append(img_name)
    
    # 为每个类别选择样本
    samples = {}
    for label, images in class_images.items():
        if len(images) >= num_per_class:
            samples[label] = random.sample(images, num_per_class)
        else:
            samples[label] = images
    
    # 创建图像网格
    num_classes = len(samples)
    fig, axes = plt.subplots(num_classes, num_per_class, figsize=(num_per_class*4, num_classes*4))
    
    # 类别名称映射
    class_names = {
        0: "类别0",
        1: "类别1",
        2: "类别2",
        3: "类别3",
        4: "类别4",
        5: "类别5"
    }
    
    # 显示图像
    for i, (label, img_names) in enumerate(sorted(samples.items())):
        for j, img_name in enumerate(img_names):
            img_path = os.path.join(image_dir, img_name)
            img = cv2.imread(img_path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            
            if num_classes == 1:
                ax = axes[j]
            else:
                ax = axes[i, j]
                
            ax.imshow(img)
            ax.set_title(f"{class_names.get(label, f'类别{label}')}")
            ax.axis('off')
    
    plt.tight_layout()
    plt.savefig('类别样本图像.png', dpi=300, bbox_inches='tight')
    plt.show()

# 分析类别不平衡性
def analyze_class_imbalance(train_counts, val_counts):
    # 合并训练集和验证集的计数
    total_counts = {}
    for cls in set(list(train_counts.keys()) + list(val_counts.keys())):
        total_counts[cls] = train_counts.get(cls, 0) + val_counts.get(cls, 0)
    
    # 计算每个类别的比例
    total_samples = sum(total_counts.values())
    class_ratios = {cls: count/total_samples for cls, count in total_counts.items()}
    
    # 创建饼图
    plt.figure(figsize=(10, 10))
    # 按照样本数量从大到小排序类别
    sorted_classes = sorted(class_ratios.keys(), key=lambda x: total_counts[x], reverse=True)
    labels = [f"类别{cls} ({class_ratios[cls]:.1%})" for cls in sorted_classes]
    sizes = [total_counts[cls] for cls in sorted_classes]
    
    # 使用不同的颜色
    colors = plt.cm.tab10(np.arange(len(sizes)) % 10)
    
    # 突出显示最大和最小的类别
    explode = [0.1 if (cls == max(total_counts, key=total_counts.get) or 
                      cls == min(total_counts, key=total_counts.get)) 
              else 0 for cls in sorted(class_ratios.keys())]
    
    plt.pie(sizes, explode=explode, labels=labels, colors=colors, autopct='%1.1f%%',
            shadow=True, startangle=140, textprops={'fontsize': 12})
    plt.axis('equal')  # 保持饼图为圆形
    plt.title('数据集类别分布比例', fontsize=16)
    plt.tight_layout()
    plt.savefig('类别比例.png', dpi=300, bbox_inches='tight')
    plt.show()

# 主函数
def main():
    print("开始分析数据集...")
    
    # 统计类别分布
    train_counts = count_classes(train_labels)
    val_counts = count_classes(val_labels)
    
    print("\n训练集类别分布(按样本数量从大到小排序):")
    for cls, count in sorted(train_counts.items(), key=lambda x: x[1], reverse=True):
        print(f"类别 {cls}: {count} 样本 ({count/len(train_labels):.1%})")
    
    print("\n验证集类别分布(按样本数量从大到小排序):")
    for cls, count in sorted(val_counts.items(), key=lambda x: x[1], reverse=True):
        print(f"类别 {cls}: {count} 样本 ({count/len(val_labels):.1%})")
    
    # 可视化类别分布
    plot_class_distribution(train_counts, val_counts)
    
    # 分析类别不平衡性
    analyze_class_imbalance(train_counts, val_counts)
    
    # 获取图像尺寸分布
    print("\n分析图像尺寸分布...")
    image_sizes = get_image_sizes(train_image_dir)
    
    # 可视化图像尺寸分布
    plot_image_size_distribution(image_sizes)
    
    # 显示每个类别的样本图像
    print("\n显示每个类别的样本图像...")
    show_sample_images(train_image_dir, train_labels)
    
    print("\n数据集分析完成!")

if __name__ == "__main__":
    main()