🔄 CL-RTDETR-DIFFUSION

Class-incremental object detection on COCO, where the replay buffer is regenerated by a diffusion model instead of stored.

❓ Key question: how do we generate new replay data while still keeping the ground-truth class and bounding box valid? — ControlNet, conditioned on a Canny edge map of the original image, keeps geometry fixed while redrawing everything else, so the original boxes still apply to the generated image.

🏆 Results

COCO 2017, two tasks, RT-DETR R50-vd @ 640×640. Baselines from their respective papers, best per column in bold.

40–40 split — best in the table on 5 of 6 metrics.

70–10 split — overall AP trails VLM-PL slightly, but 🚀 wins on AP_S and AP_L — the size extremes where old-class detail is hardest to retain.

Pseudo-labels are the dominant factor behind these gains, worth ~24 AP on 40–40.

🔬 Method

RT-DETR (R50-vd) learns COCO 2017 in disjoint groups of classes, one task at a time. Naively fine-tuning on a new task collapses accuracy on old classes: every old-class object is still in the image but has no label, so the matcher trains the model to treat it as background. We counter this with three mechanisms, active from task 1 onward:

1. 🏷️ Pseudo-labels — the frozen previous-task model relabels old classes in new-task images, so the matcher optimizes old and new jointly instead of erasing old classes.
2. 🎯 Attention distillation — MSE between old and new model's self-attention in the last encoder layer, a softer constraint with no confidence threshold.
3. 🎨 Diffusion replay — instead of storing old images verbatim, a sample of them is regenerated by Stable Diffusion + ControlNet, conditioned on a Canny edge map, so geometry (and bounding boxes) survive but texture/lighting/background are redrawn:

old image ──► Canny edges ──────┐
                                 ├──► ControlNet + SD ──► new pixels, same layout, original boxes
COCO caption + object counts ──►┘

Prompt construction for diffusion replay:

Positive prompt = the image's real COCO caption + its object counts (e.g. two dogs, one bicycle) + fixed quality words (realistic, detailed, 8K, ...). Negative prompt is a fixed list (deformed, cgi, cartoon, low quality, ...). Object counts matter because the edge map fixes where objects are but the text encoder still under-generates repeated instances.

🧩 How the ControlNet adapter is trained:

Stable Diffusion itself stays frozen; only a ControlNet adapter is trained. For each buffer image, the pipeline takes the Canny edge map + the generated prompt, feeds both into the (trainable) adapter, which conditions the (frozen) pretrained SD U-Net to predict an output image. The predicted image is compared against the original buffer image with an L2 loss, so the adapter learns to keep the diffusion output faithful to the original layout while SD's own weights are never touched.

📊 Why the buffer is filled per class, not by sampling images uniformly:

COCO's annotation counts span about three orders of magnitude — person alone has ~262k annotations while many tail classes (hair drier, toaster, scissors) have only a few hundred. Sampling buffer images uniformly at random would return a buffer dominated by person/car/chair, so buffer selection instead targets a fixed rate per class to keep rare classes represented.

🚀 Install & run

conda create -n tan python=3.10 && conda activate tan
pip install -r requirements.txt          # torch 2.0.1 / torchvision 0.15.2 / diffusers 0.31.0

python scripts/train.py                                    # task 0, default config
torchrun --nproc_per_node=<N> scripts/train.py              # multi-GPU
python scripts/train.py -c <cfg> -r ckpt.pth --test-only    # evaluate only

Before running: replace the hardcoded /workspace/... paths in configs/ with your own, and set export WANDB_MODE=offline if you don't have a W&B account (logging is otherwise mandatory).

There is no loop over tasks — one process trains one task, and task_idx must be set consistently across configs/rtdetr/include/dataloader.yml, configs/rtdetr/include/rtdetr_r50vd.yml, and configs/cl_pipeline.yml. For task 1, also point teacher_path at the task-0 checkpoint and enable pseudo_label/distill_attn.