import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms, models
from PIL import Image
import matplotlib.pyplot as plt
from sklearn.metrics import classification_report, confusion_matrix
import seaborn as sns
from tqdm import tqdm
import json
# 设置随机种子
torch.manual_seed(42)
np.random.seed(42)
# 类别定义
CLASS_NAMES = {0: 'Dent', 1: 'Hole', 2: 'Rusty'}
NUM_CLASSES = 3
class YOLOClassificationDataset(Dataset):
"""从YOLO格式标签中提取主要类别的数据集"""
def __init__(self, image_dir, label_dir, transform=None):
self.image_dir = image_dir
self.label_dir = label_dir
self.transform = transform
# 获取所有图片文件
self.image_files = [f for f in os.listdir(image_dir)
if f.endswith(('.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'))]
self.data = []
self._load_data()
def _load_data(self):
"""加载数据并提取主要类别"""
print(f"Loading data from {self.image_dir}...")
for img_file in tqdm(self.image_files, desc="Processing images"):
# 构建标签文件路径
base_name = os.path.splitext(img_file)[0]
label_file = os.path.join(self.label_dir, base_name + '.txt')
if not os.path.exists(label_file):
print(f"Warning: Label file not found for {img_file}")
continue
# 读取YOLO标签
with open(label_file, 'r') as f:
lines = f.readlines()
if len(lines) == 0:
print(f"Warning: Empty label file for {img_file}")
continue
# 提取所有类别
classes = []
for line in lines:
parts = line.strip().split()
if len(parts) >= 5:
class_id = int(parts[0])
if class_id in [0, 1, 2]: # 只保留有效类别
classes.append(class_id)
if len(classes) == 0:
print(f"Warning: No valid classes in {img_file}")
continue
# 获取主要类别（出现最多的类别）
main_class = max(set(classes), key=classes.count)
self.data.append({
'image_path': os.path.join(self.image_dir, img_file),
'label': main_class,
'all_classes': classes
})
print(f"Loaded {len(self.data)} images")
# 打印类别分布
class_counts = {0: 0, 1: 0, 2: 0}
for item in self.data:
class_counts[item['label']] += 1
print("\nClass distribution:")
for class_id, count in class_counts.items():
print(f" {CLASS_NAMES[class_id]} ({class_id}): {count} images")
def __len__(self):
return len(self.data)
def __getitem__(self, idx):
item = self.data[idx]
# 加载图片
image = Image.open(item['image_path']).convert('RGB')
# 应用变换
if self.transform:
image = self.transform(image)
label = item['label']
return image, label
def get_transforms(is_training=True):
"""获取数据增强和预处理transforms"""
if is_training:
return transforms.Compose([
transforms.Resize((256, 256)),
transforms.RandomCrop(224),
transforms.RandomHorizontalFlip(),
transforms.RandomRotation(15),
transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
else:
return transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
def create_model(num_classes=3, pretrained=True):
"""创建ResNet50分类模型"""
model = models.resnet50(pretrained=pretrained)
# 修改最后的全连接层
num_features = model.fc.in_features
model.fc = nn.Linear(num_features, num_classes)
return model
def train_epoch(model, dataloader, criterion, optimizer, device):
"""训练一个epoch"""
model.train()
running_loss = 0.0
correct = 0
total = 0
progress_bar = tqdm(dataloader, desc='Training')
for images, labels in progress_bar:
images, labels = images.to(device), labels.to(device)
# 前向传播
optimizer.zero_grad()
outputs = model(images)
loss = criterion(outputs, labels)
# 反向传播
loss.backward()
optimizer.step()
# 统计
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
# 更新进度条
progress_bar.set_postfix({
'loss': f'{loss.item():.4f}',
'acc': f'{100.*correct/total:.2f}%'
})
epoch_loss = running_loss / len(dataloader)
epoch_acc = 100. * correct / total
return epoch_loss, epoch_acc
def validate(model, dataloader, criterion, device):
"""验证模型"""
model.eval()
running_loss = 0.0
correct = 0
total = 0
all_predictions = []
all_labels = []
with torch.no_grad():
progress_bar = tqdm(dataloader, desc='Validation')
for images, labels in progress_bar:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
running_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
all_predictions.extend(predicted.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
progress_bar.set_postfix({
'loss': f'{loss.item():.4f}',
'acc': f'{100.*correct/total:.2f}%'
})
epoch_loss = running_loss / len(dataloader)
epoch_acc = 100. * correct / total
return epoch_loss, epoch_acc, all_predictions, all_labels
def plot_training_history(train_losses, train_accs, val_losses, val_accs, save_path):
"""绘制训练历史"""
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
# 损失曲线
ax1.plot(train_losses, label='Train Loss', marker='o')
ax1.plot(val_losses, label='Val Loss', marker='s')
ax1.set_xlabel('Epoch')
ax1.set_ylabel('Loss')
ax1.set_title('Training and Validation Loss')
ax1.legend()
ax1.grid(True)
# 准确率曲线
ax2.plot(train_accs, label='Train Acc', marker='o')
ax2.plot(val_accs, label='Val Acc', marker='s')
ax2.set_xlabel('Epoch')
ax2.set_ylabel('Accuracy (%)')
ax2.set_title('Training and Validation Accuracy')
ax2.legend()
ax2.grid(True)
plt.tight_layout()
plt.savefig(save_path, dpi=300, bbox_inches='tight')
print(f"Training history plot saved to {save_path}")
plt.close()
def plot_confusion_matrix(y_true, y_pred, save_path):
"""绘制混淆矩阵"""
cm = confusion_matrix(y_true, y_pred)
plt.figure(figsize=(10, 8))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
xticklabels=[CLASS_NAMES[i] for i in range(NUM_CLASSES)],
yticklabels=[CLASS_NAMES[i] for i in range(NUM_CLASSES)])
plt.title('Confusion Matrix')
plt.ylabel('True Label')
plt.xlabel('Predicted Label')
plt.tight_layout()
plt.savefig(save_path, dpi=300, bbox_inches='tight')
print(f"Confusion matrix saved to {save_path}")
plt.close()
def main():
# ================== 配置参数 ==================
TRAIN_IMAGE_DIR =
TRAIN_LABEL_DIR =
TEST_IMAGE_DIR =
TEST_LABEL_DIR =
OUTPUT_DIR =
os.makedirs(OUTPUT_DIR, exist_ok=True)
BATCH_SIZE = 32
NUM_EPOCHS = 50
LEARNING_RATE = 0.001
NUM_WORKERS = 4
# ================== 设备配置 ==================
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"\n{'='*60}")
print(f"Using device: {device}")
if torch.cuda.is_available():
print(f"GPU: {torch.cuda.get_device_name(0)}")
print(f"GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.2f} GB")
print(f"{'='*60}\n")
# ================== 加载数据 ==================
print("Loading datasets...")
train_dataset = YOLOClassificationDataset(
TRAIN_IMAGE_DIR,
TRAIN_LABEL_DIR,
transform=get_transforms(is_training=True)
)
test_dataset = YOLOClassificationDataset(
TEST_IMAGE_DIR,
TEST_LABEL_DIR,
transform=get_transforms(is_training=False)
)
train_loader = DataLoader(
train_dataset,
batch_size=BATCH_SIZE,
shuffle=True,
num_workers=NUM_WORKERS,
pin_memory=True
)
test_loader = DataLoader(
test_dataset,
batch_size=BATCH_SIZE,
shuffle=False,
num_workers=NUM_WORKERS,
pin_memory=True
)
print(f"\nTrain batches: {len(train_loader)}")
print(f"Test batches: {len(test_loader)}\n")
# ================== 创建模型 ==================
print("Creating model...")
model = create_model(num_classes=NUM_CLASSES, pretrained=True)
model = model.to(device)
# 损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
optimizer, mode='min', factor=0.5, patience=5
)
# ================== 训练模型 ==================
print("\nStarting training...\n")
train_losses = []
train_accs = []
val_losses = []
val_accs = []
best_val_acc = 0.0
best_epoch = 0
for epoch in range(NUM_EPOCHS):
print(f"\n{'='*60}")
print(f"Epoch {epoch+1}/{NUM_EPOCHS}")
print(f"{'='*60}")
# 训练
train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, device)
train_losses.append(train_loss)
train_accs.append(train_acc)
# 验证
val_loss, val_acc, val_preds, val_labels = validate(model, test_loader, criterion, device)
val_losses.append(val_loss)
val_accs.append(val_acc)
# 学习率调整
scheduler.step(val_loss)
# 打印结果
print(f"\nEpoch {epoch+1} Results:")
print(f" Train Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
print(f" Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%")
print(f" Learning Rate: {optimizer.param_groups[0]['lr']:.6f}")
# 保存最佳模型
if val_acc > best_val_acc:
best_val_acc = val_acc
best_epoch = epoch + 1
torch.save({
'epoch': epoch + 1,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'val_acc': val_acc,
'val_loss': val_loss,
}, os.path.join(OUTPUT_DIR, 'best_model.pth'))
print(f" ✓ Best model saved! (Val Acc: {val_acc:.2f}%)")
# ================== 保存最终模型 ==================
torch.save({
'epoch': NUM_EPOCHS,
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'train_losses': train_losses,
'train_accs': train_accs,
'val_losses': val_losses,
'val_accs': val_accs,
}, os.path.join(OUTPUT_DIR, 'final_model.pth'))
print(f"\nFinal model saved!")
# ================== 绘制训练曲线 ==================
plot_training_history(
train_losses, train_accs, val_losses, val_accs,
os.path.join(OUTPUT_DIR, 'training_history.png')
)
# ================== 最终评估 ==================
print(f"\n{'='*60}")
print("Final Evaluation on Test Set")
print(f"{'='*60}\n")
# 加载最佳模型
checkpoint = torch.load(os.path.join(OUTPUT_DIR, 'best_model.pth'))
model.load_state_dict(checkpoint['model_state_dict'])
# 评估
test_loss, test_acc, test_preds, test_labels = validate(
model, test_loader, criterion, device
)
print(f"\n{'='*60}")
print("Best Model Performance:")
print(f"{'='*60}")
print(f"Best Epoch: {best_epoch}")
print(f"Test Loss: {test_loss:.4f}")
print(f"Test Accuracy: {test_acc:.2f}%")
print(f"{'='*60}\n")
# 分类报告
print("Classification Report:")
print("="*60)
report = classification_report(
test_labels, test_preds,
target_names=[CLASS_NAMES[i] for i in range(NUM_CLASSES)],
digits=4
)
print(report)
# 保存分类报告
with open(os.path.join(OUTPUT_DIR, 'classification_report.txt'), 'w') as f:
f.write(f"Best Model Performance\n")
f.write(f"{'='*60}\n")
f.write(f"Best Epoch: {best_epoch}\n")
f.write(f"Test Loss: {test_loss:.4f}\n")
f.write(f"Test Accuracy: {test_acc:.2f}%\n")
f.write(f"{'='*60}\n\n")
f.write("Classification Report:\n")
f.write(report)
# 绘制混淆矩阵
plot_confusion_matrix(
test_labels, test_preds,
os.path.join(OUTPUT_DIR, 'confusion_matrix.png')
)
# 保存预测结果
predictions_data = {
'true_labels': [int(x) for x in test_labels],
'predicted_labels': [int(x) for x in test_preds],
'class_names': CLASS_NAMES
}
with open(os.path.join(OUTPUT_DIR, 'predictions.json'), 'w') as f:
json.dump(predictions_data, f, indent=2)
print(f"\n{'='*60}")
print("Training completed!")
print(f"All results saved to: {OUTPUT_DIR}")
print(f"{'='*60}\n")
# 保存训练配置
config = {
'model': 'ResNet50',
'num_classes': NUM_CLASSES,
'class_names': CLASS_NAMES,
'batch_size': BATCH_SIZE,
'num_epochs': NUM_EPOCHS,
'learning_rate': LEARNING_RATE,
'best_epoch': best_epoch,
'best_val_acc': float(best_val_acc),
'final_test_acc': float(test_acc),
'train_images': len(train_dataset),
'test_images': len(test_dataset)
}
with open(os.path.join(OUTPUT_DIR, 'config.json'), 'w') as f:
json.dump(config, f, indent=2)
if __name__ == '__main__':
main()
最新发布
在运行使用Keras的代码时遇到了KeyError: 'val_acc'。错误源于Keras版本差异，导致预期的指标名称不一致，可能是'val_acc'，'acc'或'accuracy'。解决方法是检查并适配Keras版本或修改代码以匹配正确指标名。
最低0.47元/天 解锁文章
扫一扫
3662
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
