WeedRepFormer: Reparameterizable Vision Transformers for Real-Time Waterhemp Segmentation and Gender Classification

TL;DR

WeedRepFormer tackles the real-time joint task of waterhemp segmentation and gender classification by applying comprehensive structural reparameterization to a Vision Transformer architecture, enabling high accuracy with low latency on edge devices. The model couples a reparameterizable four-stage ViT backbone (RepPatchEmbed, RepCPE, RepMixFFN), a Lite R-ASPP decoder, and a reparameterizable classification head (RepClsHead) to fuse training-time capacity into deployment efficiency. A 10,264-frame waterhemp dataset from 23 plants with pixel-accurate masks and plant-level gender labels underpins the evaluation, where WeedRepFormer achieves 92.18% mIoU and 81.91% mAcc at 108.95 FPS with only 3.59M parameters and 3.80 GFLOPs, outperforming state-of-the-art methods in classification while remaining segmentation-competitive. The approach demonstrates strong practical impact for real-time, targeted weed management by reliably identifying female waterhemp to reduce seed production and resistance evolution in agricultural settings.

Abstract

We present WeedRepFormer, a lightweight multi-task Vision Transformer designed for simultaneous waterhemp segmentation and gender classification. Existing agricultural models often struggle to balance the fine-grained feature extraction required for biological attribute classification with the efficiency needed for real-time deployment. To address this, WeedRepFormer systematically integrates structural reparameterization across the entire architecture - comprising a Vision Transformer backbone, a Lite R-ASPP decoder, and a novel reparameterizable classification head - to decouple training-time capacity from inference-time latency. We also introduce a comprehensive waterhemp dataset containing 10,264 annotated frames from 23 plants. On this benchmark, WeedRepFormer achieves 92.18% mIoU for segmentation and 81.91% accuracy for gender classification using only 3.59M parameters and 3.80 GFLOPs. At 108.95 FPS, our model outperforms the state-of-the-art iFormer-T by 4.40% in classification accuracy while maintaining competitive segmentation performance and significantly reducing parameter count by 1.9x.

Figures (5)

• Figure 1: Overview of our multi-task architecture. (a) Network consists of four-stage hierarchical Vision Transformer backbone with reparameterizable components. (b) Reparameterizable patch embedding with multi-branch convolutions. (c) LRASPP decoder with reparameterizable convolutions. (d) Classification head with optional SE attention.
• Figure 2: Reparameterizable components with train-time overparameterization. \ref{['fig:reppatchembed']} RepPatchEmbed uses three parallel branches that fuse at inference. \ref{['fig:repmixffn']} RepMixFFN employs K parallel depthwise convolutions with identity connection. \ref{['fig:repcpe']} RepCPE combines depthwise convolution with identity for positional encoding.
• Figure 3: Reparameterizable convolution modules used in decoder and classifier. \ref{['fig:repconv1x1']} RepConv1x1 with K parallel 1×1 branches for efficient channel mixing. \ref{['fig:repdwconv3x3']} RepDWConv3x3 with K parallel 3×3 depthwise branches for spatial feature refinement.
• Figure 4: Example image-mask pairs from the waterhemp dataset. Left pair shows a female specimen, right pair shows a male specimen. Binary segmentation masks delineate complete plant architecture against background.
• Figure 5: Qualitative comparison on male and female waterhemp. From left to right: input, ground truth, RepViT, FastViT, MobileOne, MobileNetV2, MobileNetV3, SegFormer-B0, and WeedRepFormer (ours). Best viewed on screen.