Feature Request: RankSEG Post-Processing Integration for Metric-Aware Inference Optimization

[Leev1s]

问题确认 Search before asking

• 我已经查询历史issue(包含open与closed)，没有类似需求。I have searched the open and closed issues and found no similar feature requests.

需求描述 Feature Description

Feature Request: RankSEG Post-Processing Integration for Metric-Aware Inference Optimization

📋 Issue Type

• Bug Report
• Feature Request
• Documentation
• Other

---

🎯 Summary

We propose to integrate RankSEG as an optional post-processing layer in PaddleSeg's inference pipeline. RankSEG is a plug-and-play, training-free module that replaces the standard argmax step with a metric-aware optimization, yielding measurable gains in Dice/IoU without any retraining.

On CamVid with PP-LiteSeg-T, we demonstrate:

Model	mIoU ($IoU^D$)	Dice	Note
PP-LiteSeg-T (baseline)	75.92%	81.38%	argmax
PP-LiteSeg-T + RankSEG	76.13%	82.48%	RankSEG

+0.21 pp mIoU, +1.10 pp Dice — zero retraining.

A fully reproducible environment is available as a pre-built Docker image: ghcr.io/leev1s/rankseg:paddleseg. See the README for exact commands.

---

🔬 Background: What is RankSEG?

RankSEG is a post-processing optimization technique published in two peer-reviewed venues:

• 📄 JMLR 2023: RankDice / RankSEG
• 📄 NeurIPS 2025: RankSEG-RMA

Core Idea

Standard segmentation pipelines end with:

pred = argmax(logit)  # maximizes Accuracy, not Dice/IoU

argmax is theoretically optimal for pixel-wise Accuracy but not for non-decomposable metrics like Dice or IoU. RankSEG bridges this gap:

# RankSEG (3 lines in pure PyTorch)
from rankseg import RankSEG
rankseg = RankSEG(metric='dice')   # or 'iou', 'acc'
pred = rankseg.predict(softmax(logit, dim=1))

Available solvers: RMA (fast, recommended), BA, TRNA, BA+TRNA.

---

⚠️ Current Limitation: The Dual-Framework Challenge

RankSEG is implemented in PyTorch. PaddleSeg is built on PaddlePaddle. Running both frameworks simultaneously in a single Python process comes with real costs.

Our current implementation explicitly acknowledges this by performing a logit.numpy() round-trip:

# rankseg/paddleseg/infer.py (our prototype)
import torch
from rankseg import RankSEG

logit_np = logit.numpy()                             # Paddle → NumPy
logit_torch = torch.from_numpy(logit_np)              # NumPy → PyTorch
probs = torch.softmax(logit_torch, dim=1)
pred_np = RankSEG(metric='dice').predict(probs).numpy()  # PyTorch → NumPy
pred = paddle.to_tensor(pred_np, dtype='int32')       # NumPy → Paddle

Our Migration Intent

We fully intend to port RankSEG's core algorithms to PaddlePaddle so that the logit never leaves the Paddle ecosystem. The RMA solver, for example, requires only basic tensor operations (sum, cumsum, sort, element-wise multiply) that map directly to paddle equivalents. We would like to coordinate with the PaddleSeg team on this transition.

The current dual-framework implementation should be viewed as a functional prototype for demonstration and benchmarking purposes, not a production design.

---

🛠️ Proposed Integration Architecture

Minimal-Invasive Approach: _generate_pred_from_logit

Rather than modifying the public API of inference() / aug_inference(), we propose inserting a small internal helper into paddleseg/core/infer.py:

def _generate_pred_from_logit(logit,
                               use_multilabel=False,
                               use_rankseg=False,
                               rankseg_metric='dice',
                               rankseg_solver='RMA',
                               rankseg_output_mode='multiclass'):
    """
    Replace argmax with a metric-aware predictor (RankSEG) when requested.
    Falls back to original PaddleSeg logic when use_rankseg=False.
    """
    if not use_rankseg:
        # ---- ORIGINAL PADDLESEG LOGIC (unchanged) ----
        if not use_multilabel:
            return paddle.argmax(logit, axis=1, keepdim=True, dtype='int32')
        else:
            return (F.sigmoid(logit) > 0.5).astype('int32')
    else:
        # ---- RANKSEG PATH ----
        # Phase 1 (current): Paddle → NumPy → PyTorch bridge
        # Phase 2 (planned): native Paddle RankSEG implementation
        import torch
        from rankseg import RankSEG

        probs = torch.softmax(torch.from_numpy(logit.numpy()), dim=1)
        rs = RankSEG(metric=rankseg_metric, solver=rankseg_solver,
                     output_mode=rankseg_output_mode)
        pred_np = rs.predict(probs).cpu().numpy()
        if rankseg_output_mode == 'multiclass' and pred_np.ndim == 3:
            pred_np = pred_np[:, None, :, :]
        return paddle.to_tensor(pred_np, dtype='int32')

Both inference() and aug_inference() call _generate_pred_from_logit(logit, ...) at their tail, passing RankSEG flags through **kwargs. The existing pred, _ = infer.inference(...) call signature is fully preserved for all callers that do not opt in.

Config & CLI Support

# paddleseg config YAML (optional new section)
post_processing:
  use_rankseg: True
  rankseg_metric: dice       # dice | iou | acc
  rankseg_solver: RMA        # RMA | BA | TRNA | BA+TRNA
  rankseg_output_mode: multiclass   # multiclass | multilabel

# tools/val.py
python tools/val.py \
  --config configs/pp_liteseg/pp_liteseg_stdc1_camvid_960x720_10k.yml \
  --model_path model.pdparams \
  --use_rankseg --rankseg_metric="iou"

---

📦 Reproducible Demo via Docker

README We provide a fully containerized environment that demonstrates the integration end-to-end.

Image Architecture

ghcr.io/leev1s/rankseg:paddleseg
├── Stage 1 (builder) — uv:python3.10-bookworm-slim
│   ├── paddlepaddle==3.3.0  (CPU build, ~700 MB)
│   └── torch (CPU build via pytorch.org/whl/cpu, ~500 MB)
├── Stage 2 (pioneer) — uv:python3.10-bookworm
│   ├── PaddleSeg release/2.10 (cloned & installed from source)
│   └── rankseg (installed from source)
└── Stage 3 (runtime) — python:3.10-slim-bookworm
    ├── /opt/venv  ← copied virtual environment
    ├── /workspace/configs  ← PaddleSeg upstream configs (release/2.10)
    └── /workspace/{val,predict,analyse}.py  ← RankSEG-aware tools

Both paddlepaddle and torch are installed in the same virtual environment. GPU version is also tested but not published.

Quick Reproduction

# Pull image
docker pull ghcr.io/leev1s/rankseg:paddleseg

# Prepare host directories
mkdir -p ./data ./models ./output
curl -L https://paddleseg.bj.bcebos.com/dataset/camvid.tar | tar -xC ./data
mkdir -p ./models/pp_liteseg_camvid
curl -L https://paddleseg.bj.bcebos.com/dygraph/camvid/pp_liteseg_stdc1_camvid_960x720_10k/model.pdparams \
     -o ./models/pp_liteseg_camvid/model.pdparams

# Run container
docker run --rm -it \
  -v ./data:/workspace/data \
  -v ./models:/workspace/pretrained_models \
  -v ./output:/workspace/output \
  ghcr.io/leev1s/rankseg:paddleseg

# Inside container:
# Baseline
python val.py \
  --config configs/pp_liteseg/pp_liteseg_stdc1_camvid_960x720_10k.yml \
  --model_path pretrained_models/pp_liteseg_camvid/model.pdparams

# With RankSEG
python val.py \
  --config configs/pp_liteseg/pp_liteseg_stdc1_camvid_960x720_10k.yml \
  --model_path pretrained_models/pp_liteseg_camvid/model.pdparams \
  --use_rankseg --rankseg_metric='iou'

---

📈 Benchmark Results

All results use $IoU^D$ (dataset-level accumulation) — the metric PaddleSeg reports by default:

$$IoU^D = \frac{\sum_{i=1}^{N} TP_i}{\sum_{i=1}^{N} (TP_i + FP_i + FN_i)}$$

$IoU^I$ (image-level average) is computed differently — it first calculates IoU for each image independently, then takes the mean:

$$IoU^I = \frac{1}{N} \sum_{i=1}^{N} \frac{TP_i}{TP_i + FP_i + FN_i}$$

Key difference: $IoU^D$ accumulates TP/FP/FN across all images before computing the ratio, while $IoU^I$ computes the ratio per-image first, then averages. $IoU^D$ is more sensitive to image size distribution (larger images contribute more to the accumulation), while $IoU^I$ treats all images equally regardless of size.

Note: For detailed discussions on the theoretical and practical differences between $IoU^D$ and $IoU^I$, please refer to the NeurIPS 2025 paper.

CamVid — PP-LiteSeg-T (STDC1)

# Baseline (argmax)
[EVAL] #Images: 101  mIoU: 0.7592  Acc: 0.9652  Kappa: 0.9564  Dice: 0.8138
[EVAL] Image-level average mIoU (IoU^I): 0.7361

# + RankSEG (metric=iou, solver=RMA)
[EVAL] #Images: 101  mIoU: 0.7622  Acc: 0.9604  Kappa: 0.9506  Dice: 0.8255
[EVAL] Image-level average mIoU (IoU^I): 0.7377

Metric	Baseline	+ RankSEG	Δ
mIoU ($IoU^D$)	75.92%	76.13%	+0.21 pp
Dice	81.38%	82.48%	+1.10 pp
mIoU ($IoU^I$)	73.61%	73.77%	+0.16 pp

🖼️ Visualization Results

Below are qualitative comparisons between baseline (argmax) and RankSEG predictions on CamVid test images:

Example 1: 0016E5_08051

Baseline (argmax)	RankSEG (iou, RMA)

Example 2: 0016E5_08159

Baseline (argmax)	RankSEG (iou, RMA)

Observations: RankSEG improves boundary delineation and reduces small misclassified regions, particularly in challenging areas like thin structures and class boundaries.

---

🔗 References

Resource	Link
RankSEG GitHub	https://github.com/rankseg/rankseg
RankSEG PyPI	https://pypi.org/project/rankseg/
RankSEG Documentation	https://rankseg.readthedocs.io/en/latest/
JMLR 2023 Paper	https://www.jmlr.org/papers/v24/22-0712.html
NeurIPS 2025 Paper	https://openreview.net/pdf?id=4tRMm1JJhw
Docker Image	ghcr.io/leev1s/rankseg:paddleseg
Integration Source	https://github.com/Leev1s/rankseg/tree/paddleseg/rankseg/paddleseg
Prototype infer.py	https://github.com/Leev1s/rankseg/blob/paddleseg/rankseg/paddleseg/infer.py

---

📝 Checklist

• Searched existing issues for duplicates
• Provided performance benchmarks
• Acknowledged the dual-framework limitation and proposed a migration path
• Provided a fully reproducible Docker environment

Special thanks to @statmlben @ZixunWang

是否愿意提交PR Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!