Custom Helmet Detection Dataset

Dataset Overview

Introduction

This dataset is used to train a custom object‑detection model (based on TensorFlow 2.x) that automatically detects whether a person is wearing a helmet. The dataset originates from a custom helmet‑detection collection on Kaggle.

Dataset Preparation

1. Load the dataset: It contains 764 images divided into two classes for helmet detection. Bounding‑box annotations use the PASCAL VOC format.
2. Annotate bounding boxes: If the dataset lacks bounding boxes, they must be labeled manually. A common tool is labelImg.
3. Create a label‑map file (.pbtxt): A label‑map must be defined for each class.

   item {
       id: 1
       name: With Helmet
   }
   
   item {
       id: 2
       name: Without Helmet
   }
   

4. Create TFRecord files (.record): Convert the dataset to TFRecord format for training the custom detector.

Pre‑trained Model

The project uses the SSD MobileNet V2 FPNLite 320x320 pre‑trained model for transfer learning.

Training

Before training, modify the pipeline.config file to set the number of classes, batch size, pre‑trained model path, etc.

Training Command

python model_main_tf2.py --model_dir=[model_directory_path] --pipeline_config_path=[pipeline_config_path]

Output

Sample output images after training illustrate the helmet‑detection results.