Skip to content

Model returns class IDs outside the dataset category range during inference #930

@Isabelam00

Description

@Isabelam00

Search before asking

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

Bug

I trained a model using the following dataset:

[(https://universe.roboflow.com/wyhil-ru2ds/workers-safety-equipment-z1mra/dataset/3)]

The dataset contains 8 categories with IDs from 0 to 7. However, during inference the model sometimes returns class IDs outside this range (for example 73).

Dataset categories:

"categories": [
{"id": 0, "name": "worker"},
{"id": 1, "name": "Gloves"},
{"id": 2, "name": "Helmet"},
{"id": 3, "name": "Non-Helmet"},
{"id": 4, "name": "Person"},
{"id": 5, "name": "Shoes"},
{"id": 6, "name": "Vest"},
{"id": 7, "name": "bare-arms"}
]

Expected behavior:

The model should only return class IDs within the dataset range (0–7).

Actual behavior:

Sometimes the model returns class IDs outside this range (e.g., 73).

Question:

Is this expected behavior (for example due to pretraining on a larger dataset like COCO), or could there be an issue with how the checkpoint was trained or loaded?

Environment

  • RFDETR: 1.4.1
  • Python 3.13.3

Minimal Reproducible Example

Train script

from rfdetr import RFDETRMedium

model = RFDETRMedium()

model.train(
    	dataset_dir="/Datasets train/workers-safety-equipment.v3i.coco",
    	epochs=20,
    	batch_size=4,
    	grad_accum_steps=4,
    	lr=1e-4,
    	output_dir="/Epi_workers-safety-equipment_Medium_v2")

Inference script

import cv2
import supervision as sv
from rfdetr import RFDETRMedium
import json
import os

coco_json_path = "test/_annotations.coco.json"

with open(coco_json_path, "r") as f:
    data = json.load(f)

id_to_class = {c["id"]: c["name"] for c in data["categories"]}

model = RFDETRMedium(pretrain_weights="checkpoint_best_total_20epochs.pth")

url_video = 'ruta_video'
cap = cv2.VideoCapture(url_video)

while True:
    success, frame = cap.read()
    if not success:
        break    

    detections = model.predict(frame, threshold=0.5)
    print('classes:', detections.class_id)
    
    labels = [
        f"{id_to_class.get(class_id, 'unknown')} {confidence:.2f}"
        for class_id, confidence in zip(detections.class_id, detections.confidence)
    ]

    annotated_frame = frame.copy()
    annotated_frame = sv.BoxAnnotator().annotate(annotated_frame, detections)
    annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)

    cv2.imshow("Webcam", annotated_frame)

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

Example outputs:

classes: [ 3 4 2 73]
classes: [3 4 2]
classes: [3 4 2]
classes: [4 1 3 2 1]

Additional

No response

Are you willing to submit a PR?

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions