FireSegUNet 🔥 — Fire Segmentation on Edge Devices

Official PyTorch implementation of the paper:
"FireSegUNet: Exploring Computationally Efficient Fire Segmentation Network for Unmanned Aerial Vehicles"
Accepted at Knowledge-Based Systems (Elsevier, ISSN 0950-7051), June 2026.

---

Overview

FireSegUNet is a lightweight, modular deep-learning architecture designed for fire segmentation on resource-constrained devices such as UAVs, IoT platforms, and embedded edge hardware.
It features:

• Optimized Inverted Bottleneck Layer (OIBL) for efficient feature extraction
• Efficient Squeeze-and-Excitation Block (ESEB) for lightweight channel attention
• 81% memory reduction with competitive segmentation accuracy
• Better inference speeds on mobile GPUs and embedded processors (ARM Cortex, Adreno, etc.)

Although developed primarily for fire segmentation, FireSegUNet generalizes well to other segmentation tasks across domains like environmental monitoring, IoT, TinyML, and mobile vision applications.

---

Features

• PyTorch + Lightning implementation with modular codebase
• Pretrained models on FLAME, BoWFire, and Fire datasets
• Training, evaluation, and ONNX export scripts
• Benchmarking utilities (FLOPs, MACs, inference time)
• Easy to customize for other segmentation tasks

---

Paper

Ali Hassan, Mårten Sjöström, Karen Egiazarian, Tingting Zhang, Johan Johansson, Stefan Schulte
FireSegUNet: Exploring Computationally Efficient Fire Segmentation Network for UAVs, 2025.

Paper Link (Coming Soon)

---

Project Structure

FireSegUNet/
├── main.py                   # Training & inference script
├── evaluate.py               # Evaluation & ONNX export script
├── eval.py                   # Lightweight evaluation script
├── class_accuracy.py         # Per-class accuracy computation
├── requirements.txt          # Python dependencies
├── utils/
│   ├── firenet.py            # FireNet wrapper (SMP encoders)
│   ├── firenet_baseline.py   # FireNet wrapper (custom architectures)
│   ├── flamedataset_loader.py
│   ├── bowdataset_loader.py
│   ├── firedataset_loader.py
│   ├── all_dataloader.py     # Combined dataset loader
│   ├── dataset_visualizer.py
│   └── model_library/        # Custom model architectures
│       ├── unet_proposed.py  # FireSegUNet (proposed)
│       ├── dsnet.py
│       ├── manet.py
│       ├── unet.py / unet_ibn.py / unet_2plus.py / unet_3plus.py
│       └── ...
├── segmentation_models_pytorch/   # Modified SMP library
├── lightning_logs/               # Training checkpoints & TensorBoard logs
└── tests/                        # Unit tests

---

Installation

Step 1 — Install PyTorch with CUDA support first:

pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu117

Step 2 — Install segmentation-models-pytorch:

pip install -U segmentation-models-pytorch

Step 3 — Install remaining dependencies:

pip install -r requirements.txt

Note: The segmentation_models_pytorch/ directory in this repo contains a modified version of SMP used internally. The pip install above provides the base library; local modifications override it at runtime via the included package.

---

Quick Start

Training

Edit main.py to select dataset and model, then run:

# In main.py:
dataType = 1   # 1: FLAME  2: BoWFire  3: FireDataset  4: Mixed
model_name = 'FireSegUNet'
initChanel = 32

python main.py

Evaluation

Edit evaluate.py to select the dataset and checkpoint path, then run:

python evaluate.py

Outputs: IoU, F1, Precision, Recall, Sensitivity, Specificity, and an exported .onnx model.

Supported Architectures

Model	Description
FireSegUNet	Proposed lightweight model (OIBL + ESEB)
DSNet / DSNet-Base	Dual-stream network
UNet / AttUNet / R2UNet / R2AttUNet	UNet variants
NestedUNet / UNet2Plus / UNet3Plus	Nested UNet variants
MANet	Multi-scale Attention Network
DeepLabv3+	Optimized DeepLab variant
MixUNet	Mixed-architecture UNet

Supported Encoders (SMP-based models)

mobilenet_v2, mobileone_s0, mit_b0, vgg11, timm-mobilenetv3_small_100,
timm-efficientnet-b0, resnet18, densenet121, inceptionv4, xception

---

Pretrained Checkpoints

Pretrained Lightning checkpoints are provided under lightning_logs/:

Dataset	Checkpoint Path
FLAME	lightning_logs/flame/FLAME-UNet-Proposed2-SE-RELU6-UpConv-C32/version_0/checkpoints/
BoWFire	lightning_logs/bow/BoW-UNet-Proposed2-SE-RELU6-UpConv-C32/version_0/checkpoints/
FireDataset	lightning_logs/fire/Fire-UNet-Proposed2-SE-RELU6-UpConv-C32/version_0/checkpoints/
Mixed	lightning_logs/mixed/UNet-Proposed2-RELU6-SE-C32/version_0/checkpoints/

---

Benchmarks

Dataset	Mean IoU (%)	F1 Score (%)	Parameters (M)	Memory Reduction
FLAME	85.25	92.04	1.28	−81%

Full benchmarking details available in the paper (link coming soon).

---

Code Release Timeline

Stage	Status
Code Cleaning	Done
Paper Submission	Accepted
Repository Public Release	Coming Soon upon publication

---

Acknowledgments

This work was supported by the European Joint Doctoral Programme on Plenoptic Imaging (PLENOPTIMA) through the European Union's Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie Grant Agreement No. 956770. The research was also supported by the European Union through the EU Interreg Aurora project IMMERSE (20366448), and by Mid Sweden University internal funding. The computations and data handling were enabled by resources provided by the National Academic Infrastructure for Supercomputing in Sweden (NAISS) at Chalmers Centre for Computational Science and Engineering (C3SE), partially funded by the Swedish Research Council through grant agreement no. 2022-06725.

---

Citation

If you use FireSegUNet or this code, please cite:

@article{hassan2026firesegunet,
  title={FireSegUNet: Exploring Computationally Efficient Fire Segmentation Network for Unmanned Aerial Vehicles},
  author={Hassan, Ali and Sj{\"o}str{\"o}m, M{\aa}rten and Egiazarian, Karen and Zhang, Tingting and Johansson, Johan and Schulte, Stefan},
  journal={Knowledge-Based Systems},
  issn={0950-7051},
  publisher={Elsevier},
  month={june},
  year={2026}
}