Datawhale AI春训营 + 智慧骑士—消防隐患识别

消防隐患识别

本文记录了我在楼道消防隐患识别任务中，将目标检测与图像分类融合的完整实践过程，包括模型架构设计、训练策略以及最终评估结果。

---

背景与目标

对“消防隐患随手拍”项目中的拍摄照片内容进行识别，实时判断照片内场景是否存在消防安全隐患以及隐患的危险程度。根据楼道中是否存在大量堆积物，堆积物是否可燃以及是否有起火风险将隐患分为无隐患、低风险、中等风险、高风险。

高风险：楼道中出现电动车、电瓶、飞线充电等可能起火的元素。

中风险：楼道中存在大量堆积物严重影响通行或堆放大量纸箱、木质家具等能造成火势蔓延的堵塞物。

低风险：存在楼道堆物现象但不严重。

无风险：楼道干净，无堆放物品。

非楼道：一些与楼道无关的图片。

细则：

1、高风险场景需要有过道中停放电动自行车、给电瓶充电、楼道中飞线充电等一项行为或多项行为。

2、中风险场景需要至少满足以下两个条件之一：①楼道内堆积众多堆积物已经严重影响通行。②楼道的堆积物中有明显可见的纸箱、木质或布质的家具、泡沫箱等可燃物品。

3、低风险场景主要为楼道中有物品摆放但基本不影响通行，数量较少或靠边有序摆放。

---

评估与测量:

w_1=2, w_2=1.5,w_3=w_4=w_5=1

模型架构设计

Baseline

Baseline 仅仅只是一个十分简单的ResNet18的模型

def get_model1():
    model = models.resnet18(True)
    model.fc = nn.Linear(512, 5)
    return model

跑出来的score仅有 2.0左右。
所以我们需要设计更多新的模型用来提高准确率

1. 模型介绍：ResNet50 + EfficientNet3b

最简单的一个思路就是用更大的ResNet 融合 更大的EfficientNet来做5分类
所以我们采取一个简单的双分支融合思路

ResNet50分支

def get_resnet():
    model = models.resnet50(pretrained=True)
    model.fc = nn.Linear(2048, 5)  # 修改最后一层
    return model

EfficientNet-B3分支

def get_efficientnet():
    model = models.efficientnet_b3(pretrained=True)
    model.classifier = nn.Sequential(
        nn.Dropout(0.5),
        nn.Linear(1536, 128),  # 中间特征层
        nn.ReLU(),
        nn.Linear(128, 5)
  )
    return model

融合分支

class EnsembleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.resnet = get_resnet()       # ResNet50主干
        self.effnet = get_efficientnet() # EfficientNet-B3主干
        self.classifier = nn.Sequential(
            nn.Linear(10, 128),          # 特征融合层
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(128, 5)           # 最终分类层
        )

def forward(self, x):
    x1 = self.resnet(x)    # ResNet特征 [B,5]
    x2 = self.effnet(x)    # EfficientNet特征 [B,5]
    x = torch.cat([x1, x2], dim=1)  # 拼接特征 [B,10]
    return self.classifier(x)  # 融合分类 [B,5]

代码流程

数据预处理

def prepare_data():
    train_df = pd.read_csv(Config.train_data_path, sep="\t", header=None)
    train_df[0] = Config.image_dir + train_df[0]
    
    # 过滤无效图片
    def is_valid_image(image_path):
        if not os.path.exists(image_path):
            return False
        img = cv2.imread(image_path)
        return img is not None
    
    train_df = train_df[train_df[0].apply(is_valid_image)]
    
    # 标签映射
    mapping_dict = {'高风险':0, '中风险':1, '低风险':2, '无风险':3, '非楼道':4}
    train_df[1] = train_df[1].map(mapping_dict)
    class_counts = train_df[1].value_counts().sort_index().values
    class_weights = 1. / torch.tensor(class_counts, dtype=torch.float)
    sample_weights = class_weights[train_df[1].values]
    sampler = torch.utils.data.WeightedRandomSampler(
        sample_weights, len(sample_weights), replacement=True
    )
    
    return train_df

读取图片，过滤无效图片

数据增强

train_transform = transforms.Compose([
        transforms.Resize((Config.img_size, Config.img_size)),
        transforms.RandomHorizontalFlip(),
        transforms.RandomVerticalFlip(),
        transforms.ColorJitter(0.2, 0.2, 0.2),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    val_transform = transforms.Compose([
        transforms.Resize((Config.img_size, Config.img_size)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

损失函数构建

因为数据集中标签不均衡的问题，我们需要构建一个新的权重损失函数

class DynamicWeightedFocalLoss(nn.Module):
    def __init__(self, base_alpha, gamma=2):
        self.base_alpha = torch.tensor([2.0,1.5,1.0,1.0,1.0])
        self.gamma = gamma  # 困难样本聚焦参数

    def forward(self, inputs, targets):
        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
        pt = torch.exp(-ce_loss)  # 分类置信度
        focal_loss = (1-pt)**self.gamma * ce_loss
        return (self.base_alpha[targets] * focal_loss).mean()

训练过程

def train_epoch(model, loader, criterion, optimizer, scaler=None):
    model.train()
    total_loss = 0.0
    correct = 0
    
    for inputs, targets in loader:
        inputs, targets = inputs.to(Config.device), targets.to(Config.device)
        optimizer.zero_grad()
        
        # 混合精度上下文
        with torch.autocast(device_type=Config.device.type, enabled=Config.use_amp):
            outputs = model(inputs)
            loss = criterion(outputs, targets)
        
        if scaler:
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            scaler.update()
        else:
            loss.backward()
            optimizer.step()
        
        total_loss += loss.item()
        preds = outputs.argmax(dim=1)
        correct += (preds == targets).sum().item()
    
    return total_loss/len(loader), correct/len(loader.dataset)

def validate(model, loader, criterion):
    model.eval()
    total_loss = 0.0
    correct = 0
    
    with torch.no_grad():
        for inputs, targets in loader:
            inputs, targets = inputs.to(Config.device), targets.to(Config.device)
            
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            total_loss += loss.item()
            preds = outputs.argmax(dim=1)
            correct += (preds == targets).sum().item()
    
    return total_loss/len(loader), correct/len(loader.dataset)

数据可视化

def plot_history(history, fold):
    epochs = list(range(1, len(history['train_loss'])+1))
    plt.figure()
    plt.plot(epochs, history['train_loss'], label='Train Loss')
    plt.plot(epochs, history['val_loss'], label='Val Loss')
    plt.title(f'Fold {fold} Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.legend()
    plt.savefig(f"Loss1_{fold}.jpg")

    plt.figure()
    plt.plot(epochs, history['train_acc'], label='Train Acc')
    plt.plot(epochs, history['val_acc'], label='Val Acc')
    plt.title(f'Fold {fold} Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy')
    plt.legend()
    plt.savefig(f'Loss2_{fold}.jpg')

预测

test_df = pd.read_csv(Config.test_data_path, sep="\t", header=None)
    test_df["path"] = Config.test_dir + test_df[0]
    
    test_transform = transforms.Compose([
        transforms.Resize((Config.img_size, Config.img_size)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    test_loader = DataLoader(
        CustomDataset(test_df["path"].values, 
        np.zeros(len(test_df)),  # dummy labels
        transform=test_transform
    ), batch_size=Config.batch_size*2, 
       shuffle=False,
       num_workers=Config.num_workers,
       pin_memory=(Config.device.type == "cuda"))
    
    # 加载最佳模型
    probs_sum = np.zeros((len(test_df), 5), dtype=np.float32)
    for fold in range(1, 6):
        model = EnsembleModel().to(Config.device)
        if torch.cuda.device_count()>1:
            model = nn.DataParallel(model)
        model.load_state_dict(torch.load(f'best_model_fold{fold}.pth', map_location=Config.device))
        model.eval()
        with torch.no_grad():
            idx = 0
            for imgs, _ in test_loader:
                imgs = imgs.to(Config.device)
                outputs = model(imgs)
                probs = nn.functional.softmax(outputs, dim=1).cpu().numpy()
                batch_size = imgs.size(0)
                probs_sum[idx:idx+batch_size] += probs
                idx += batch_size
    
    # 保存结果
    avg_preds = probs_sum.argmax(axis=1)
    inv_map = {0:'高风险',1:'中风险',2:'低风险',3:'无风险',4:'非楼道'}
    test_df['label'] = [inv_map[p] for p in avg_preds]
    test_df[[0, 'label']].to_csv('submission.txt', index=False, header=None, sep="\t")

结果示例

score = 4.6695

总结

当然，以上模型也十分简单，主要是为大家展示一下整体的流程跑通的样子。后续有时间的话我也会给大家介绍一下怎样设计一个比较好的模型。