Datawhale 春训营第二期 楼道图像分类 task2
前言
这篇笔记是对ELE AI算法大赛“赛道二：智慧骑士—消防隐患识别”baseline的提升，具体的baseline可以看我上一篇博客，在优化baseline的时候，我运用了一些工具的辅助以及task2的资料，以下是对赛题的简要描述。
赛题背景
在履约配送过程中，蓝骑士穿梭于大街小巷及楼梯拐角，能够及时发现各类隐患，堪称城市安全隐患的“移动监测员”。然而，隐患上报链路存在两大卡点：一是蓝骑士上传信息的真实性和准确性难以保证；二是上报行为的奖励反馈不够及时，这在一定程度上影响了隐患上报的积极性。借助AI的实时识别功能，可有效解决上述问题，既能确保隐患信息真实有效，又能对蓝骑士进行即时奖励，鼓励其积极上报隐患。
赛事任务
本次比赛主题为“消防隐患随手拍”项目中的照片内容识别。需要实时判断拍摄照片内场景是否存在消防安全隐患以及隐患的危险程度。根据楼道中堆积物的情况、堆积物的可燃性以及起火风险，将隐患分为无隐患、低风险、中等风险、高风险四个等级，同时识别非楼道场景。
具体等级划分标准如下：
- 高风险：楼道中出现电动自行车停放、电瓶充电、飞线充电等可能引发火灾的行为之一或多项。
- 中风险：需满足以下两个条件之一：①楼道内堆积物众多，严重影响通行；②楼道堆积物中有明显可见的纸箱、木质家具、布质家具、泡沫箱等可燃物品。
- 低风险：楼道中有物品摆放，但基本不影响通行，数量较少且靠边有序摆放。
- 无风险：楼道干净整洁，无堆放物品。
- 非楼道：与楼道无关的图片。
参赛者需依据上述标准，精准识别照片中的场景，为消防安全隐患的及时发现与处理提供技术支持。
细则：
1、高风险场景需要有过道中停放电动自行车、给电瓶充电、楼道中飞线充电等一项行为或多项行为。
2、中风险场景需要至少满足以下两个条件之一：①楼道内堆积众多堆积物已经严重影响通行。②楼道的堆积物中有明显可见的纸箱、木质或布质的家具、泡沫箱等可燃物品。
3、低风险场景主要为楼道中有物品摆放但基本不影响通行，数量较少或靠边有序摆放。
初赛评分标准
数据格式
优化后的baseline
import os
import time
import numpy as np
import pandas as pd
from PIL import Image
from datetime import datetime
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models, transforms
from sklearn.model_selection import StratifiedKFold
import albumentations as A
from albumentations.pytorch import ToTensorV2
# ===================== 设置与准备 =====================
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 标签映射与反映射
mapping_dict = {'高风险': 0, '中风险': 1, '低风险': 2, '无风险': 3, '非楼道': 4}
inverse_mapping_dict = {v: k for k, v in mapping_dict.items()}
# ===================== 数据处理类 =====================
from PIL import UnidentifiedImageError
class GalaxyDataset(Dataset):
def __init__(self, image_paths, labels, transform=None):
self.image_paths = image_paths
self.labels = labels
self.transform = transform
def __getitem__(self, index):
try:
image = np.array(Image.open(self.image_paths[index]).convert("RGB")) # 转成 numpy 数组
except (UnidentifiedImageError, OSError) as e:
print(f"Warning: Skipping corrupted image: {self.image_paths[index]}")
return self.__getitem__((index + 1) % len(self.image_paths)) # 循环避开坏图
if self.transform:
augmented = self.transform(image=image)
image = augmented["image"]
label = self.labels[index]
return image, torch.tensor(label, dtype=torch.long)
def __len__(self):
return len(self.labels)
# ===================== 模型构建 =====================
def get_model():
model = models.efficientnet_b0(pretrained=True)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 5)
return model.to(device)
# ===================== Focal Loss =====================
class FocalLoss(nn.Module):
def __init__(self, gamma=2.0, weight=None):
super().__init__()
self.gamma = gamma
self.weight = weight
def forward(self, input, target):
ce_loss = F.cross_entropy(input, target, reduction='none', weight=self.weight)
pt = torch.exp(-ce_loss)
return ((1 - pt) ** self.gamma * ce_loss).mean()
# ===================== 训练与验证 =====================
def train_one_epoch(model, loader, criterion, optimizer):
model.train()
running_loss = 0
for images, labels in loader:
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() * images.size(0)
return running_loss / len(loader.dataset)
def validate(model, loader):
model.eval()
correct = 0
with torch.no_grad():
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
correct += (outputs.argmax(1) == labels).sum().item()
return correct / len(loader.dataset)
# ===================== 训练入口 =====================
def train_cv():
df = pd.read_csv("data/train.txt", sep="\t", header=None)
df[0] = "data/train/" + df[0]
df = df[df[0].apply(lambda x: os.path.exists(x))]
df[1] = df[1].map(mapping_dict)
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
for fold, (train_idx, val_idx) in enumerate(skf.split(df[0], df[1])):
print(f"===== Fold {fold + 1} =====")
train_paths, val_paths = df.iloc[train_idx][0].values, df.iloc[val_idx][0].values
train_labels, val_labels = df.iloc[train_idx][1].values, df.iloc[val_idx][1].values
train_transform = A.Compose([
A.Resize(256, 256),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.ColorJitter(p=0.3),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2()
])
val_transform = A.Compose([
A.Resize(256, 256),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2()
])
train_loader = DataLoader(GalaxyDataset(train_paths, train_labels, train_transform),
batch_size=32, shuffle=True, num_workers=4)
val_loader = DataLoader(GalaxyDataset(val_paths, val_labels, val_transform),
batch_size=32, shuffle=False, num_workers=4)
model = get_model()
criterion = FocalLoss()
optimizer = optim.AdamW(model.parameters(), lr=3e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
best_acc = 0
for epoch in range(10):
print(f"Epoch {epoch + 1}/10")
loss = train_one_epoch(model, train_loader, criterion, optimizer)
acc = validate(model, val_loader)
scheduler.step()
print(f"Train Loss: {loss:.4f} | Val Acc: {acc:.4f}")
if acc > best_acc:
best_acc = acc
torch.save(model.state_dict(), "best_model.pt")
break # 只用第一个 Fold
# ===================== 推理入口 =====================
def predict():
df = pd.read_csv("data/A.txt", sep="\t", header=None)
df["path"] = "data/A/" + df[0]
df["label"] = 0
test_transform = A.Compose([
A.Resize(256, 256),
A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
ToTensorV2()
])
test_loader = DataLoader(GalaxyDataset(df["path"].values, df["label"].values, test_transform),
batch_size=32, shuffle=False, num_workers=4)
model = get_model()
model.load_state_dict(torch.load("best_model.pt"))
model.eval()
predictions = []
with torch.no_grad():
for images, _ in test_loader:
images = images.to(device)
outputs = model(images)
preds = outputs.argmax(1).cpu().numpy()
predictions.extend(preds)
df["label"] = [inverse_mapping_dict[p] for p in predictions]
df[[0, "label"]].to_csv("submit.txt", index=False, sep="\t", header=None)
# ===================== 主函数入口 =====================
if __name__ == "__main__":
train_cv()
predict()
提交结果
可以清楚地看到上大分了！
优化了哪些方面
模型结构
-
EfficientNet-B0是一类轻量高效的卷积网络，结构更深、更宽，且引入了 Swish 激活、MBConv 块、更好的参数利用率； -
修改最后全连接层为
nn.Linear(in_features, 5)，适配5类分类任务； -
使用
model.to(device)自动兼容 GPU/CPU。
def get_model():
model = models.efficientnet_b0(pretrained=True)
model.classifier[1] = nn.Linear(model.classifier[1].in_features, 5)
return model.to(device)
数据增强
Albumentations提供高效且丰富的图像增强方式，对图像内容和结构有更强的保留能力。增强多样性有助于模型泛化和鲁棒性提升，增加了更丰富的增强手段：HorizontalFlip, VerticalFlipColorJitter（模拟光线变化）Resize, Normalize, ToTensorV2
import albumentations as A
from albumentations.pytorch import ToTensorV2
train_transform = A.Compose([
A.Resize(224, 224),
A.HorizontalFlip(p=0.5),
A.VerticalFlip(p=0.5),
A.ColorJitter(brightness=0.2, contrast=0.2),
A.Normalize(),
ToTensorV2(),
])
错误图像处理
在 Dataset.__getitem__ 中加入 try/except：捕获 UnidentifiedImageError, OSError；遇到异常自动跳过并使用下一个图像，防止训练中断，提升训练稳定性和自动化程度。
def __getitem__(self, index):
try:
image = np.array(Image.open(self.image_paths[index]).convert("RGB")) # 转成 numpy 数组
except (UnidentifiedImageError, OSError) as e:
print(f"Warning: Skipping corrupted image: {self.image_paths[index]}")
return self.__getitem__((index + 1) % len(self.image_paths)) # 循环避开坏图
if self.transform:
augmented = self.transform(image=image)
image = augmented["image"]
label = self.labels[index]
return image, torch.tensor(label, dtype=torch.long)
损失函数改进：引入 Focal Loss
Focal Loss 是为应对类别不均衡而设计的，尤其适用于小样本类别在多分类任务中被忽视的情况，实现了 FocalLoss，加入 γ 调节因子，样本困难度越高，loss 权重越大，鼓励模型关注 hard examples；
class FocalLoss(nn.Module):
def __init__(self, gamma=2.0):
super().__init__()
self.gamma = gamma
self.ce = nn.CrossEntropyLoss()
def forward(self, input, target):
logp = self.ce(input, target)
p = torch.exp(-logp)
loss = (1 - p) ** self.gamma * logp
return loss.mean()
criterion = FocalLoss(gamma=2.0)
训练调度器
- 使用带权重衰减的
Adam - CosineAnnealing 可提升收敛稳定性和泛化能力，引入 CosineAnnealingLR 调度器，动态衰减学习率至最低点再回升，有助于跳出局部最优；
- 初始学习率设置为 3e-4
optimizer = optim.AdamW(model.parameters(), lr=3e-4)
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
best_acc = 0
for epoch in range(10):
print(f"Epoch {epoch + 1}/10")
loss = train_one_epoch(model, train_loader, criterion, optimizer)
acc = validate(model, val_loader)
scheduler.step()
print(f"Train Loss: {loss:.4f} | Val Acc: {acc:.4f}")
if acc > best_acc:
best_acc = acc
torch.save(model.state_dict(), "best_model.pt")
break # 只用第一个 Fold
优化推理
- 使用 .argmax(dim=1) 判断预测类别；
- .sum().item() 统计预测正确数；
- 除以样本总数得到准确率
test_loader = DataLoader(GalaxyDataset(df["path"].values, df["label"].values, test_transform),
batch_size=32, shuffle=False, num_workers=4)
model = get_model()
model.load_state_dict(torch.load("best_model.pt"))
model.eval()
predictions = []
with torch.no_grad():
for images, _ in test_loader:
images = images.to(device)
outputs = model(images)
preds = outputs.argmax(1).cpu().numpy()
predictions.extend(preds)
个人心得
在优化baseline的时候，我们可以从多个关键维度上进行系统性改进，来提升模型泛化性能、训练稳定性以及鲁棒性，我在提供的能跑通的baseline的基础上，利用task02资料的辅助，以及在AI的帮助下，成功使用了一些模型，在这个过程中不仅学会了如何使用某一些模型，模型的原理，还有上分的原因，可能在使用模型，以及理解的时候比较费时间，但是收获满满啊~
扫一扫
688
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
