How to run object detection inference with a YOLOE segmentation model?

Discussion
1

How to load YOLOE model without the segmentation head

YOLOE support training for object detection, and I am trying to verify inference of an YOLOE object detection models works and benchmark it

model = YOLO("yoloe-11l.yaml").load('panopticon_models/yoloe-11l-seg.pt')
model.set_classes(["person"])   # this fails 
model.predict()

Note: I am not up-to-date with how the latest YOLO architectures do segmentation, the latest paper I have read is YOLOACT, and I assume that YOLO{NUMEBER}-seg models work in a similar way (by extending detection head to detect mask)

Thanks!

2

The code you provided works fine for me. What exactly “fails”?

3 
$ py
Python 3.12.11 (main, Jul 23 2025, 11:34:34) [GCC 15.1.1 20250425] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from ultralytics import YOLO
>>> model = YOLO("yoloe-11l.yaml").load('panopticon_models/yoloe-11l-seg.pt')
Transferred 1108/1108 items from pretrained weights
>>> model.set_classes(["person"])
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/master-andreas/panopticon/test_env/lib/python3.12/site-packages/ultralytics/models/yolo/model.py", line 330, in set_classes
    embeddings = self.get_text_pe(classes)  # generate text embeddings if not provided
                 ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/master-andreas/panopticon/test_env/lib/python3.12/site-packages/ultralytics/models/yolo/model.py", line 269, in get_text_pe
    return self.model.get_text_pe(texts)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/master-andreas/panopticon/test_env/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/master-andreas/panopticon/test_env/lib/python3.12/site-packages/ultralytics/nn/tasks.py", line 1038, in get_text_pe
    assert not self.training
           ^^^^^^^^^^^^^^^^^
AssertionError
4

You’re using set_classes() incorrectly.

5

that is a perfectly valid way of using YOLOE.set_classes, but even if the more verbose way of using the function is used (as shown in the documentation), it fails in the same way

from ultralytics import YOLO

model = YOLO("yoloe-11l.yaml").load('panopticon_models/yoloe-11l-seg.pt')
names = ['person']
model.set_classes(names, model.get_text_pe(names))  # this also fails
model.set_classes(names)  # this fails

$ py test.py                                    
Transferred 1108/1108 items from pretrained weights
Traceback (most recent call last):
  File "/home/master-andreas/panopticon/test.py", line 5, in <module>
    model.set_classes(names, model.get_text_pe(names))  # this also fails
                             ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/master-andreas/panopticon/test_env/lib/python3.12/site-packages/ultralytics/models/yolo/model.py", line 269, in get_text_pe
    return self.model.get_text_pe(texts)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/master-andreas/panopticon/test_env/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 120, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/master-andreas/panopticon/test_env/lib/python3.12/site-packages/ultralytics/nn/tasks.py", line 1038, in get_text_pe
    assert not self.training
           ^^^^^^^^^^^^^^^^^
AssertionError
6

Seems like your transferring of the weights is causing the issue.

You can do model.eval() before set_classes()

1 Like
7

It works thanks

from ultralytics import YOLO

model = YOLO("yoloe-11s.yaml").load('panopticon_models/yoloe-11s-seg.pt')
model.eval()
names = ['person']
model.set_classes(names, model.get_text_pe(names))  # this works 
model.set_classes(names)  # this also works

What does YOLOE.eval() do, I can not seem to find any details on the documentation, (I can only find the val() method, which does something different)

8

That is because it’s not an Ultralytics specific method, it’s inherited from the PyTorch Module class and you can check those docs here

1 Like