Skip to content

Automatic labeling tool for image segmentation based on detection-labeled dataset. Additionally provides the tools for fine tuning a smoke segmenter.

License

Notifications You must be signed in to change notification settings

fl0wxr/SmokeSegmenter

Repository files navigation

Smoke Segmentation Project

Smoke segmentation data is scarce online, and the already existing available datasets are either poorly segmented or limited in size. The objective of this repository is to facilitate the creation of an expansive image dataset specifically tailored for smoke segmentation tasks, leveraging a semi-automatic labeling approach based on a foundation model called SAM. This repository offers a comprehensive suite of tools, notably including a segmentation UI, designed for an oracle to filter out (or blacklist) poorly predicted pseudo-ground-truth segmentation masks. Moreover, it encompasses methodologies and tools essential for training smoke segmentation models.

Fine-Tuned Models

Two models have been trained on our smoke segmentation dataset to get a comparison of how effective each is on the task of smoke segmentation.

Results

Model Pre-trained @ mIU mIU Smoke FPS
BiSeNet - R18 Cityscapes 80.81% 69.48% 21.4
PIDNet Small Camvid 81.64% 70.69% 25.0

Details - Demo

The resulting predictions along with a metrics diagram follows, showcasing the BiSeNet (R18 backbone) model's performance and its training potential.



Figure 1. The left image shows the input, the middle image is the combination of the input with the predicted segmentation mask, and the right image is the (pseudo) ground truth mask acquired by SAM.


Figure 2. Depicted mean IU plot computed on the test set.

Specs & OS

All the experiments referenced were conducted on a system equipped with an Intel i5-12500H CPU, 38.9 GB of memory, and an RTX 4060 GPU (laptop version), running Ubuntu 22.04.3 LTS.

D-Fire - External Dataset

[D-Fire] is an image dataset of fire and smoke occurrences designed for machine learning and object detection algorithms with more than 21,000 images. The bounding box labels are stored inside .txt files in YOLO format (class + cxcywh).

S-Smoke

We name S-Smoke to be any dataset consisted of D-Fire's images, with segmentation ground truth masks produced for the task of image smoke segmentation, where the ground truth masks are generated from a pre-trained [SAM] model, prompted by the bounding box labels of the D-Fire dataset. An oracle then filters out bad predictions.

Our produced dataset is named as S-Smoke-var0, and is consisted by

  • Number of train instances 2400
  • Number of test instances 287

Data Directory Structure

This is the data's directory structure, located in ./datasets.

datasets
├── data_info
├── D-Fire
|   ├── test
|   |   ├── images
|   |   └── det_labels
|   └── train
|       ├── images
|       └── det_labels
└── S-Smoke
    ├── curated
    |   ├── test
    |   |   ├── images
    |   |   ├── seg_labels
    |   |   └── combined
    |   └── train
    └── raw
        ├── test
        |   ├── images
        |   ├── seg_labels
        |   └── combined
        └── train

Model Directory Structure

This directory structure has to be built manually.

models
├── PIDNet
|   ├── finetuned
|   └── pretrained
|       └── PIDNet_S_Camvid_Test.pt
└── TorchSeg
    ├── finetuned
    └── pretrained
        └── cityscapes-bisenet-R18.pth

Download and place these at their corresponding directories

Preparation

Directory

The developer has to manually create the preceeding directory structures prior to using any of the provided tools. The D-Fire directory has to be filled with the corresponding files from the dataset. D-Fire's labels and images are expected to be in text and JPG file formats, and be named with the corresponding .txt and .jpg file extensions. Download and decompress the D-Fire dataset.

This is what the decompressed D-Fire directory will look like initially:

D-Fire
├── test
|   ├── det_labels
|   └── images
└── train
    ├── det_labels
    └── images

To align the paths with the currently provided scripts, you should rename the directories with names labels to det_labels.

Now download a pretrained SAM model through [link] and place it inside ./models directory.

Dependencies

Apply

sudo apt-get install ninja-build
python3 -m pip install -r requirements.txt

Finally install [Apex].

Next Steps

First visit [data_tools] to perform the semi-annotation by carefully following all instructions given by the corresponding README.md. After that, you will acquire your own S-Smoke dataset and you may train your own segmentation models through [primary_segmenter].

Citation