Segmentation issue with line "model.optimize_for_inference()"

[hangxu811]

Search before asking

• [x] I have searched the RF-DETR issues and found no similar bug report.

Bug

I fine-tune rf-detr seg preview for my project, and I found out that when I deployed the model and do the inference without this line "model.optimize_for_inference()", everything works fine like below

detections = model.predict(image, threshold=0.3) print(detections)

Detections(xyxy=array([[888.2473 , 746.7094 , 932.44073, 768.10944]], dtype=float32), mask=array([[[False, False, False, ..., False, False, False], [False, False, False, ..., False, False, False], [False, False, False, ..., False, False, False], ..., [False, False, False, ..., False, False, False], [False, False, False, ..., False, False, False], [False, False, False, ..., False, False, False]]], shape=(1, 1024, 1280)), confidence=array([0.75640106], dtype=float32), class_id=array([1]), tracker_id=None, data={}, metadata={})

but when I implement this line of code "model.optimize_for_inference()", the model stops generating mask.

Detections(xyxy=array([[888.24963, 746.70496, 932.4463 , 768.10974]], dtype=float32), mask=None, confidence=array([0.75613964], dtype=float32), class_id=array([1]), tracker_id=None, data={}, metadata={})

For people who want to train segmentation, please do not use this line of code as they use in their repo predict example.

Environment

RF-DETR 1.3

Minimal Reproducible Example

custom_weights_path = "/root/runs/rfdetr_seg/checkpoint_best_regular.pth"

try: model = RFDETRSegPreview(pretrain_weights=custom_weights_path) except Exception as e: model = RFDETRSegPreview()

detections = model.predict(image, threshold=0.3)

print(detections)

Additional

No response

Are you willing to submit a PR?

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

[isaacrob-roboflow]

Does it work when using the most recent source code?

[J4BEZ]

@isaacrob-roboflow I believe I've found a fix for this issue. May I open a pull request to propose the change?

[isaacrob-roboflow]

@J4BEZ i believe it is already working in the raw code, just not in the release version

[hangxu811]

@J4BEZ i believe it is already working in the raw code, just not in the release version

Thank you so much!

[isaacrob-roboflow]

Can I make this as completed @hangxu811 ?

[hangxu811]

Yes please

[J4BEZ]

@isaacrob-roboflow @hangxu811

I apologize for the delayed response.

I encountered the same issue when building from source, so I’m sharing a minimal reproducible example here for reference.

minimal code representation

# Create env for representation
# uv venv ./.envs --python 3.11
# source ./.envs/bin/activate

# install rf-detr from raw code
!uv pip install torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu128
!git clone https://github.com/roboflow/rf-detr.git
%cd rf-detr
!uv pip install -e .

!uv pip show rfdetr
# Using Python 3.11.13 environment at: .../workspace/.envs
# Name: rfdetr
# Version: 1.3.0
# Location .../workspace/.envs/lib/python3.11/site-packages
# Editable project location: .../workspace/rf-detr
# Requires: accelerate, cython, einops, fairscale, matplotlib, ninja, numpy, open-clip-torch, pandas, peft, polygraphy, pycocotools, pydantic, pylabel, rf100vl, roboflow, scipy, supervision, timm, torch, torchvision, tqdm, transformers
# Required-by:

import supervision as sv
from supervision.annotators.base import BaseAnnotator

from rfdetr import RFDETRSegPreview
from rfdetr.util.coco_classes import COCO_CLASSES

from PIL import Image
import requests

from typing import Tuple

# Annotator example
class Annotators:
    def __init__(self, *args):
        self.annotators: Tuple[BaseAnnotator, ...] = args
    
    def annotate(self, scene: Image.Image, detections: sv.Detections) -> Image.Image:
        scene = scene.copy()
        for annotator in self.annotators:

            # special handling for LabelAnnotator to add class name with confidence
            if issubclass(type(annotator), sv.LabelAnnotator):
                labels = [
                    f"{COCO_CLASSES[class_id]} {confidence:.2f}"
                    for class_id, confidence
                    in zip(detections.class_id, detections.confidence)
                ]
                scene = annotator.annotate(scene=scene, detections=detections, labels=labels)
            else:
                scene = annotator.annotate(scene=scene, detections=detections)
        return scene

annotators = Annotators(
    sv.BoxAnnotator(thickness=1),
    sv.LabelAnnotator(),
    sv.MaskAnnotator(),
    sv.PolygonAnnotator(thickness=1)
)


# Sample Image
url =  "https://media.roboflow.com/notebooks/examples/dog.jpeg"
image_data = requests.get(url, stream=True).raw
image = Image.open(image_data)

model = RFDETRSegPreview()
detections = model.predict(image)
print(type(detections.mask), detections.mask.shape if detections.mask is not None else "detections.mask is None")

annotated_image = annotators.annotate(scene=image, detections=detections)
annotated_image

# <class 'numpy.ndarray'> (3, 1280, 720)

model.optimize_for_inference()
detections = model.predict(image)
print(type(detections.mask), detections.mask.shape if detections.mask is not None else "detections.mask is None")

annotated_image = annotators.annotate(scene=image, detections=detections)
annotated_image

# <class 'NoneType'> detections.mask is None

model.remove_optimized_model()
detections = model.predict(image)
print(type(detections.mask), detections.mask.shape if detections.mask is not None else "detections.mask is None")

annotated_image = annotators.annotate(scene=image, detections=detections)
annotated_image

# <class 'numpy.ndarray'> (3, 1280, 720)

If possible, I would be grateful for the opportunity to contribute to the RF-DETR project. (It seems I’ve found a way to fix this issue🙌) Thank you again for your hard work on this excellent project and hope you have a great day!

[isaacrob-roboflow]

Go ahead! We're SUPER short on bandwidth so idk when we'll be able to review / merge tho, apologies

[J4BEZ]

@isaacrob-roboflow Thank you! I'm delighted to have the opportunity to contribute to rf-detr. I’ll open a PR soon!

[J4BEZ]

@isaacrob-roboflow Hello!

While I was working on preparing the PR, I noticed that another pull request addressing this exact issue has already been submitted.

• #429
• #430

In light of the solution already submitted by another contributor, I will refrain from submitting my PR to respect their effort and prevent redundant review cycles.

Thank you again for your hard work on this project, and I hope you have a peaceful day and stay healthy.

Best Regards, J4BEZ😊