AGCR-DETR

Code for AGCR-DETR, an RT-DETR-L based detector for bearing surface defects. The model combines amplitude–gradient feature gating (APDCCM-Lite) with Spatially Guarded Competitive Ranking (SGCR) supervision.

Code-only release: trained weights, pretrained initialization files, datasets, annotations, experimental predictions/results, and manuscript files are not included. Bring your own legally obtained data and initialization checkpoint.

Browse or clone the public source at WEFF12/AGCR-DETR:

git clone https://huggingface.co/WEFF12/AGCR-DETR
cd AGCR-DETR

Model

• APDCCM-Lite is applied independently to the C3 and C4 features before the original 1×1 projections (Ultralytics RT-DETR-L layers 19 and 14). It combines channel gating with a spatial gate derived from mean absolute amplitude and Sobel gradient magnitude, followed by a learnable residual scaling factor.
• SGCR adds a quality-weighted ranking loss to the final non-denoising decoder output during training. It compares the matched class with eligible inter-class and nearby unmatched-background competitors. Geometry and matching are detached; gradients flow through the selected probabilities.
• The original RT-DETR losses, matching and inference structure are retained. SGCR adds no inference parameters. APDCCM-Lite adds 165,908 parameters.

Installation

Use Python 3.10 or newer. Install a compatible PyTorch and torchvision build for your machine, then install the remaining dependencies:

python -m pip install -r requirements.txt

The archived main-model training used PyTorch 2.7.1+cu118 and Ultralytics 8.3.235. Ultralytics is pinned because the extension uses its RT-DETR graph and training interfaces. The lightweight CPU check below uses random tensors, requires no weights or data, and does not train a detector:

python smoke_check.py

Release checks passed with PyTorch 2.11.0 and Ultralytics 8.3.235 on a CPU. They check construction, forward execution and the loss gradient; they are not a new full training run or a measurement of detection accuracy.

Data

Copy configs/data.example.yaml to a local YAML file and set path to your own dataset directory. Use YOLO detection labels with zero-based class IDs and normalized class x_center y_center width height rows:

your_dataset/
  images/train/   images/val/   images/test/
  labels/train/   labels/val/   labels/test/

Use the same class order in your annotations and YAML. The example lists the four classes used by the main model; edit it to match your own data. The scripts do not download, select, split, or redistribute a dataset.

Training

Obtain a compatible upstream RT-DETR-L COCO initialization checkpoint separately. Run from this repository's root with one CUDA GPU:

python train.py --data configs/data.local.yaml --weights /path/to/rtdetr-l.pt --device 0 --name agcr-detr

configs/train.yaml contains the main-model settings: seed 42, 640-pixel input, batch size 4, AdamW, at most 100 epochs, and patience 40. The APDCCM initialization uses its original internal seed 42001 without advancing the baseline RNG stream. Training uses AMP with an initial GradScaler value of 128; validation and saved weights use FP32. The best state is selected by validation mAP50–95. This package preserves final-only checkpoint writing: best.pt and last.pt are written at normal training completion, so an interrupted run has no rolling resume file.

Outputs are saved under runs/train/. They are ignored by Git and are not part of this code release. The training entry point does not automatically evaluate the test set. Keep training, validation and final testing separate.

Evaluation and prediction

These commands require a checkpoint that you train locally. This repository does not provide the paper's trained checkpoints.

python evaluate.py --data configs/data.local.yaml --weights runs/train/agcr-detr/weights/best.pt --split test --device 0
python predict.py --weights runs/train/agcr-detr/weights/best.pt --source /path/to/images --conf 0.25 --device 0

Evaluation uses the checked RT-DETR validator with FP32, per-axis box scaling, and no added NMS. Prediction uses the upstream RT-DETR predictor and saves confidence scores. Use --device cpu for evaluation or prediction on a CPU. Only load checkpoints from sources you trust.

Repository structure

project_models/       APDCCM-Lite and SGCR implementations
tools/                Original model-specific trainer and validator integration
configs/              Training settings and dataset schema example
train.py              Standalone training entry point
evaluate.py           Validation/test entry point
predict.py            Local-image inference entry point
smoke_check.py        CPU model-wiring and loss-gradient checks

The model modules and five trainer/validator implementation files retain the archived experiment implementation. The standalone command-line entry points replace local experiment orchestration and machine-specific paths. This release does not include ablation variants, comparison models, or experiment archives.

License and acknowledgments

AGPL-3.0; see LICENSE and NOTICE.md. This implementation builds on Ultralytics RT-DETR.