重构：将 Baseline.py 迁移至 baseline/ 目录，新增多模型对比脚本
This commit is contained in:
@@ -0,0 +1,145 @@
|
||||
"""
|
||||
baseline/VGG_KNN.py
|
||||
VGG16 预训练模型特征提取 + KNN 四分类基线
|
||||
可独立运行，也可被 compare_models.py 导入复用
|
||||
author: yukun-hh
|
||||
date: 2026-5-14
|
||||
"""
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision import models, transforms
|
||||
from tqdm import tqdm
|
||||
|
||||
from sklearn.neighbors import KNeighborsClassifier
|
||||
from sklearn.metrics import (
|
||||
accuracy_score, f1_score,
|
||||
confusion_matrix, ConfusionMatrixDisplay,
|
||||
classification_report,
|
||||
)
|
||||
|
||||
from Dataloader import RobustImageFolder
|
||||
|
||||
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
|
||||
matplotlib.rcParams['axes.unicode_minus'] = False
|
||||
|
||||
|
||||
CLASS_NAMES = ['厨余垃圾', '可回收物', '其他垃圾', '有害垃圾']
|
||||
|
||||
|
||||
def load_vgg16_extractor(device):
|
||||
try:
|
||||
model = models.vgg16(weights='IMAGENET1K_V1')
|
||||
except TypeError:
|
||||
model = models.vgg16(pretrained=True)
|
||||
model.classifier = nn.Identity()
|
||||
model = model.to(device).eval()
|
||||
for param in model.parameters():
|
||||
param.requires_grad = False
|
||||
return model
|
||||
|
||||
|
||||
def extract_features(model, loader, device):
|
||||
model.eval()
|
||||
all_features = []
|
||||
all_labels = []
|
||||
with torch.no_grad():
|
||||
for images, labels in tqdm(loader, desc='Extracting features'):
|
||||
images = images.to(device)
|
||||
feats = model(images)
|
||||
all_features.append(feats.cpu().numpy())
|
||||
all_labels.append(labels.numpy())
|
||||
return np.concatenate(all_features), np.concatenate(all_labels)
|
||||
|
||||
|
||||
class VGGKNNBaseline:
|
||||
def __init__(self, k=5, device='cpu',
|
||||
data_root='../trash_division_data/ultimate_4_class/',
|
||||
image_size=256, batch_size=32, num_workers=4):
|
||||
self.k = k
|
||||
self.device = device
|
||||
self.data_root = data_root
|
||||
self.image_size = image_size
|
||||
self.batch_size = batch_size
|
||||
self.num_workers = num_workers
|
||||
self.extractor = load_vgg16_extractor(device)
|
||||
self.knn = KNeighborsClassifier(n_neighbors=k, n_jobs=-1)
|
||||
|
||||
def _get_loader(self, split):
|
||||
transform = transforms.Compose([
|
||||
transforms.Resize((self.image_size, self.image_size)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
dataset = RobustImageFolder(
|
||||
root=os.path.join(self.data_root, split),
|
||||
transform=transform,
|
||||
)
|
||||
print(f" {split}: {len(dataset)} 张")
|
||||
return DataLoader(dataset, batch_size=self.batch_size,
|
||||
shuffle=False, num_workers=self.num_workers,
|
||||
pin_memory=True, drop_last=False)
|
||||
|
||||
def fit(self, train_loader=None):
|
||||
if train_loader is None:
|
||||
train_loader = self._get_loader('train')
|
||||
print(" 提取训练集特征 ...")
|
||||
train_feats, train_labels = extract_features(self.extractor, train_loader, self.device)
|
||||
self.knn.fit(train_feats, train_labels)
|
||||
|
||||
def predict(self, val_loader=None):
|
||||
if val_loader is None:
|
||||
val_loader = self._get_loader('val')
|
||||
print(" 提取验证集特征 ...")
|
||||
val_feats, val_labels = extract_features(self.extractor, val_loader, self.device)
|
||||
preds = self.knn.predict(val_feats)
|
||||
probs = self.knn.predict_proba(val_feats)
|
||||
return val_labels, preds, probs
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
DATA_ROOT = '../trash_division_data/ultimate_4_class/'
|
||||
BATCH_SIZE = 32
|
||||
IMAGE_SIZE = 256
|
||||
NUM_WORKERS = 4
|
||||
K = 5
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available()
|
||||
else 'xpu' if hasattr(torch, 'xpu') and torch.xpu.is_available()
|
||||
else 'cpu')
|
||||
print(f"Device: {device}")
|
||||
|
||||
baseline = VGGKNNBaseline(k=K, device=device,
|
||||
data_root=DATA_ROOT, image_size=IMAGE_SIZE,
|
||||
batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)
|
||||
|
||||
train_loader = baseline._get_loader('train')
|
||||
val_loader = baseline._get_loader('val')
|
||||
|
||||
baseline.fit(train_loader)
|
||||
y_true, y_preds, y_probs = baseline.predict(val_loader)
|
||||
|
||||
acc = accuracy_score(y_true, y_preds)
|
||||
macro_f1 = f1_score(y_true, y_preds, average='macro')
|
||||
print(f"\n验证集 Accuracy: {acc:.4f}")
|
||||
print(f"验证集 Macro-F1: {macro_f1:.4f}")
|
||||
print(f"\n分类报告:\n{classification_report(y_true, y_preds, target_names=CLASS_NAMES)}")
|
||||
|
||||
cm = confusion_matrix(y_true, y_preds)
|
||||
fig, ax = plt.subplots(figsize=(8, 7))
|
||||
ConfusionMatrixDisplay(cm, display_labels=CLASS_NAMES).plot(
|
||||
ax=ax, cmap='Blues', values_format='d', xticks_rotation=30)
|
||||
ax.set_title(f'Baseline Confusion Matrix (VGG16 + KNN, K={K})', fontsize=14)
|
||||
plt.tight_layout()
|
||||
out_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'vgg_knn_confusion_matrix.png')
|
||||
plt.savefig(out_path, dpi=150, bbox_inches='tight')
|
||||
plt.show()
|
||||
print(f"混淆矩阵已保存: {out_path}")
|
||||
@@ -0,0 +1 @@
|
||||
# baseline package
|
||||
@@ -0,0 +1,180 @@
|
||||
"""
|
||||
baseline/compare_models.py
|
||||
多模型对比：ROC 曲线 + 准确率柱状图
|
||||
添加新模型只需在 MODELS 列表加一行，无需修改绘图代码
|
||||
author: yukun-hh
|
||||
date: 2026-5-14
|
||||
"""
|
||||
import sys, os
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
import numpy as np
|
||||
import matplotlib.pyplot as plt
|
||||
import matplotlib
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision import transforms
|
||||
from tqdm import tqdm
|
||||
|
||||
from sklearn.metrics import roc_curve, auc, accuracy_score
|
||||
|
||||
from Model import Net
|
||||
from Dataloader import RobustImageFolder
|
||||
from baseline.VGG_KNN import VGGKNNBaseline
|
||||
|
||||
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
|
||||
matplotlib.rcParams['axes.unicode_minus'] = False
|
||||
|
||||
# ============================================================
|
||||
# ★★★ 可配置参数 ★★★
|
||||
# ============================================================
|
||||
DATA_ROOT = '../trash_division_data/ultimate_4_class/'
|
||||
BATCH_SIZE = 32
|
||||
IMAGE_SIZE = 256
|
||||
NUM_WORKERS = 4
|
||||
K_KNN = 5
|
||||
# ============================================================
|
||||
|
||||
CLASS_NAMES = ['厨余垃圾', '可回收物', '其他垃圾', '有害垃圾']
|
||||
NUM_CLASSES = 4
|
||||
|
||||
# ============================================================
|
||||
# 预测函数 — 每个函数签名: (train_loader, val_loader, device) -> (y_true, y_preds, y_probs)
|
||||
# ============================================================
|
||||
|
||||
def get_resnet34_preds(train_loader, val_loader, device):
|
||||
model = Net(num_classes=NUM_CLASSES)
|
||||
state_dict = torch.load('best_model.pth', map_location='cpu')
|
||||
if 'model_state_dict' in state_dict:
|
||||
state_dict = state_dict['model_state_dict']
|
||||
elif 'model' in state_dict:
|
||||
state_dict = state_dict['model']
|
||||
model.load_state_dict(state_dict)
|
||||
model = model.to(device).eval()
|
||||
|
||||
y_true, y_preds, y_probs = [], [], []
|
||||
with torch.no_grad():
|
||||
for images, labels in tqdm(val_loader, desc='ResNet-34'):
|
||||
images, labels = images.to(device), labels
|
||||
logits = model(images)
|
||||
probs = torch.softmax(logits, dim=1)
|
||||
preds = probs.argmax(dim=1)
|
||||
y_true.append(labels.numpy())
|
||||
y_preds.append(preds.cpu().numpy())
|
||||
y_probs.append(probs.cpu().numpy())
|
||||
return np.concatenate(y_true), np.concatenate(y_preds), np.concatenate(y_probs)
|
||||
|
||||
|
||||
def get_vgg_knn_preds(train_loader, val_loader, device):
|
||||
baseline = VGGKNNBaseline(k=K_KNN, device=device)
|
||||
baseline.fit(train_loader)
|
||||
return baseline.predict(val_loader)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# ★ 模型注册表 — 添加新模型只需在这里加一行 ★
|
||||
# ============================================================
|
||||
|
||||
MODELS = [
|
||||
('ResNet-34', get_resnet34_preds),
|
||||
('VGG16 + KNN (K=5)', get_vgg_knn_preds),
|
||||
# 未来轻松扩展示例：
|
||||
# ('ResNet-18 (pretrained)', get_resnet18_preds),
|
||||
# ('ResNet-50 (pretrained)', get_resnet50_preds),
|
||||
# ('ResNet-34 (finetuned)', get_finetuned_preds),
|
||||
]
|
||||
|
||||
# ============================================================
|
||||
# 调色板 (扩展时无需修改)
|
||||
# ============================================================
|
||||
COLORS = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd', '#8c564b',
|
||||
'#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
|
||||
|
||||
|
||||
def compute_macro_roc(y_true, y_probs):
|
||||
one_hot = np.eye(NUM_CLASSES)[y_true]
|
||||
fpr_dict, tpr_dict = {}, {}
|
||||
for c in range(NUM_CLASSES):
|
||||
fpr_dict[c], tpr_dict[c], _ = roc_curve(one_hot[:, c], y_probs[:, c])
|
||||
all_fpr = np.unique(np.concatenate([fpr_dict[c] for c in range(NUM_CLASSES)]))
|
||||
mean_tpr = np.zeros_like(all_fpr)
|
||||
for c in range(NUM_CLASSES):
|
||||
mean_tpr += np.interp(all_fpr, fpr_dict[c], tpr_dict[c])
|
||||
mean_tpr /= NUM_CLASSES
|
||||
macro_auc = auc(all_fpr, mean_tpr)
|
||||
return all_fpr, mean_tpr, macro_auc
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
out_dir = os.path.dirname(os.path.abspath(__file__))
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available()
|
||||
else 'xpu' if hasattr(torch, 'xpu') and torch.xpu.is_available()
|
||||
else 'cpu')
|
||||
print(f"Device: {device}")
|
||||
|
||||
val_transform = transforms.Compose([
|
||||
transforms.Resize((IMAGE_SIZE, IMAGE_SIZE)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.485, 0.456, 0.406],
|
||||
std=[0.229, 0.224, 0.225]),
|
||||
])
|
||||
|
||||
train_dataset = RobustImageFolder(root=os.path.join(DATA_ROOT, 'train'),
|
||||
transform=val_transform)
|
||||
val_dataset = RobustImageFolder(root=os.path.join(DATA_ROOT, 'val'),
|
||||
transform=val_transform)
|
||||
print(f"训练集: {len(train_dataset)} 验证集: {len(val_dataset)}")
|
||||
|
||||
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False,
|
||||
num_workers=NUM_WORKERS, pin_memory=True, drop_last=False)
|
||||
val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False,
|
||||
num_workers=NUM_WORKERS, pin_memory=True, drop_last=False)
|
||||
|
||||
# ———— 评估所有模型 ————
|
||||
results = {}
|
||||
for name, func in MODELS:
|
||||
print(f"\n{'='*50}")
|
||||
print(f"评估: {name}")
|
||||
y_true, y_preds, y_probs = func(train_loader, val_loader, device)
|
||||
acc = accuracy_score(y_true, y_preds)
|
||||
fpr, tpr, roc_auc = compute_macro_roc(y_true, y_probs)
|
||||
results[name] = {'y_true': y_true, 'y_preds': y_preds, 'y_probs': y_probs,
|
||||
'acc': acc, 'fpr': fpr, 'tpr': tpr, 'auc': roc_auc}
|
||||
print(f" Accuracy: {acc:.4f} | Macro-AUC: {roc_auc:.4f}")
|
||||
|
||||
# ———— ROC 对比图 ————
|
||||
fig, ax = plt.subplots(figsize=(8, 7))
|
||||
for i, (name, r) in enumerate(results.items()):
|
||||
color = COLORS[i % len(COLORS)]
|
||||
ax.plot(r['fpr'], r['tpr'], color=color, lw=2,
|
||||
label=f"{name} (AUC={r['auc']:.4f})")
|
||||
ax.plot([0, 1], [0, 1], 'k--', lw=1, alpha=0.5)
|
||||
ax.set_xlim(0, 1); ax.set_ylim(0, 1.05)
|
||||
ax.set_xlabel('False Positive Rate'); ax.set_ylabel('True Positive Rate')
|
||||
ax.set_title('ROC Curve Comparison (Macro-Average)', fontsize=14)
|
||||
ax.legend(loc='lower right'); ax.grid(True, alpha=0.3)
|
||||
plt.tight_layout()
|
||||
roc_path = os.path.join(out_dir, 'roc_comparison.png')
|
||||
plt.savefig(roc_path, dpi=150, bbox_inches='tight')
|
||||
plt.show()
|
||||
print(f"\nROC 对比图已保存: {roc_path}")
|
||||
|
||||
# ———— 准确率柱状图 ————
|
||||
names = list(results.keys())
|
||||
accs = [results[n]['acc'] for n in names]
|
||||
fig, ax = plt.subplots(figsize=(8, 5))
|
||||
bar_colors = [COLORS[i % len(COLORS)] for i in range(len(names))]
|
||||
bars = ax.bar(names, accs, color=bar_colors, edgecolor='white', linewidth=1.2)
|
||||
for bar, acc in zip(bars, accs):
|
||||
ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.005,
|
||||
f'{acc:.4f}', ha='center', va='bottom', fontsize=12, fontweight='bold')
|
||||
ax.set_ylim(0, max(accs) * 1.15)
|
||||
ax.set_ylabel('Accuracy'); ax.set_title('Accuracy Comparison', fontsize=14)
|
||||
ax.grid(True, alpha=0.3, axis='y')
|
||||
plt.tight_layout()
|
||||
bar_path = os.path.join(out_dir, 'accuracy_bar.png')
|
||||
plt.savefig(bar_path, dpi=150, bbox_inches='tight')
|
||||
plt.show()
|
||||
print(f"准确率柱状图已保存: {bar_path}")
|
||||
Reference in New Issue
Block a user