ArmFormer: Lightweight Transformer Architecture for Real-Time Multi-Class Weapon Segmentation and Classification

TL;DR

ArmFormer tackles real-time, pixel-level multi-class weapon segmentation for security by integrating a CBAM-enhanced MixVisionTransformer with a lightweight hamburger decoder. It delivers state-of-the-art metrics ($80.64\%$ mIoU and $89.13\%$ mFscore) at $82.26$ FPS with a compact footprint ($4.886$ GFLOPs, $3.66$ M parameters), outperforming heavier models by up to $48\times$ in compute. The approach uses a four-stage CBAM-enhanced encoder and CBAM-integrated decoder with end-to-end cross-entropy optimization, reinforced by ablations showing the proposed uniform CBAM configuration yields the best accuracy-efficiency trade-off. This work enables robust, edge-deployed weapon segmentation across handgun, rifle, knife, revolver, and human classes, supporting real-time threat assessment on portable cameras, drones, and embedded AI accelerators; it also outlines practical directions for further efficiency gains and multi-sensor fusion.

Abstract

The escalating threat of weapon-related violence necessitates automated detection systems capable of pixel-level precision for accurate threat assessment in real-time security applications. Traditional weapon detection approaches rely on object detection frameworks that provide only coarse bounding box localizations, lacking the fine-grained segmentation required for comprehensive threat analysis. Furthermore, existing semantic segmentation models either sacrifice accuracy for computational efficiency or require excessive computational resources incompatible with edge deployment scenarios. This paper presents ArmFormer, a lightweight transformer-based semantic segmentation framework that strategically integrates Convolutional Block Attention Module (CBAM) with MixVisionTransformer architecture to achieve superior accuracy while maintaining computational efficiency suitable for resource-constrained edge devices. Our approach combines CBAM-enhanced encoder backbone with attention-integrated hamburger decoder to enable multi-class weapon segmentation across five categories: handgun, rifle, knife, revolver, and human. Comprehensive experiments demonstrate that ArmFormer achieves state-of-the-art performance with 80.64% mIoU and 89.13% mFscore while maintaining real-time inference at 82.26 FPS. With only 4.886G FLOPs and 3.66M parameters, ArmFormer outperforms heavyweight models requiring up to 48x more computation, establishing it as the optimal solution for deployment on portable security cameras, surveillance drones, and embedded AI accelerators in distributed security infrastructure.

Figures (4)

• Figure 1: Visualization of original images, masked annotations, and segmented results for each weapon category including handgun, human, knife, revolver, and rifle.
• Figure 2: ArmFormer Architecture Overview: The framework consists of a four-stage CBAM-enhanced MixVisionTransformer encoder backbone with progressive spatial resolution reduction and channel expansion, followed by a CBAM-integrated lightweight hamburger decoder head. The encoder processes input images through hierarchical stages (Stage 1-4) with progressive patch embeddings (P-3) and dual attention mechanisms (CBAM with Channel and Spatial Attention). Multi-scale features from all stages undergo overlap patch merging and are fused through the CBAM-MVT Encoder with Mixvision-Trans and CBAM modules. The decoder employs a squeeze convolution block, dual CBAM modules configured as a hamburger structure, and a classification head to generate pixel-level segmentation predictions for multi-class weapon detection.
• Figure 3: Qualitative comparison of segmentation results across different models on challenging weapon detection scenarios. The figure presents three representative test cases: revolver detection (top row), rifle segmentation (middle row), and handgun identification (bottom row). Each column represents predictions from different models (HrNetSunXLW19, ICNetzhao2018icnet, PspNetzhao2017pspnet, Uppernet_swinxiao2018unified, Segmenterstrudel2021segmenter, CGNetwu2020cgnet, EncNetZhang_2018_CVPR, and our ArmFormer along with the original image). Ground truth annotations and background regions are shown with corresponding color coding for each weapon category. ArmFormer (rightmost column) consistently produces the most accurate pixel-level segmentation with precise object boundaries and minimal false predictions across all weapon categories.
• Figure 4: Accuracy-efficiency trade-off visualization for ArmFormer ablation variants. The scatter plot illustrates the relationship between inference speed (FPS) and segmentation accuracy (mIoU) across different configurations. Our base ArmFormer (blue circle) achieves the optimal balance, positioned in the upper-right region with highest mIoU (80.64%) and competitive FPS (82.26). The dashed reference lines intersect at ArmFormer's performance point, demonstrating its superior positioning in the accuracy-efficiency space.