EfficientNet訓練自己的煙火識別算法

基於EfficientNet的煙火識別算法.

EfficientNet是19年提出的一種複合縮放模型,在較低複雜度上可以獲得很好的分類性能.詳情見EfficientNet(https://arxiv.org/abs/1905.11946)

EfficientNet-Pytorch(https://github.com/lukemelas/EfficientNet-PyTorch)

本文工作主要參考以上兩點,具體實現可以參考我的github(https://github.com/midasklr/FireSmokeDetectionByEfficientNet)

1. EfficientNet

很不錯的模型,創新性一般,工程指導比較有意義.後面我會給出我的EfficientDet實現.論文細節不贅述.

2. 煙火識別

2.1 idea

做目標檢測對數據集要求太高,沒有煙火的檢測數據集,所以我在網上爬取了大量煙火的圖片作爲EfficientNet的分類數據.

使用EfficientNet作爲特徵提取層,輸出全連接層修改爲3個節點進行煙火預測.

煙火檢測主要是通過將圖片分塊預測,更精細的結果可以使用不同大小滑動窗進行預測,對結果進行NMS.

2.2 網絡可視化

這裏可視化部分特徵圖

可以看出來,有些卷積核可以很好發現火的區域,對應特徵圖中白色激活區域.有的特徵圖擅長髮現其他特徵.

3. EfficientNet訓練自數據集

這裏提供了腳本train.py可以使用EfficientNet訓練自己的數據集.通過下面命令制定你的數據位置,使用網絡和目標類別數.

python train.py --data [ur dataset path] --arch [efficientnet model:efficientnet-b0-7] --num_cls [ur task class num]  

預測時記得相應的修改你的標籤映射文件.

 

數據集:鏈接: https://pan.baidu.com/s/1RkX76xJuop1Q6XB-f22WTg 提取碼: 65mv