36 num_classes model for inference, the obtained class_id is over 36

[wxqiankai]

Search before asking

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

Bug

I specified num_classes when loading the model.

model = init_model(model_path, resolution=1024, num_classes=36)
def init_model(model_path, resolution=384, num_classes=90):
    model = RFDETRNano(pretrain_weights=model_path, resolution=resolution, num_classes=num_classes)
    return model

then I check the args, show num_classes=36

print(vars(model.model))
{'args': Namespace(num_classes=36, ...

The visualization results are correct, indicating that the pre-trained model has been loaded. However, the num_classes of the model is incorrect, and class_ids exceeding 36 will occur during inference.

File "d:\projects\my-project\tools\common_tools\inference_images_create_labelme_json.py", line 40, in inference_images_and_save_labelmejson
    class_name = meta_info[class_id]['class_name']
KeyError: np.int64(56)

Environment

• RFDETR 0.1.3
• torch 2.6.0+cu118
• OS windows11
• python 3.9.23
  -GPU RTX2080

Minimal Reproducible Example

train args:
model = RFDETRNano(num_classes=36, resolution=1024, gradient_checkpointing=True)
init model:
model = init_model(model_path, resolution=1024, num_classes=36)
def init_model(model_path, resolution=384, num_classes=90):
model = RFDETRNano(pretrain_weights=model_path, resolution=resolution, num_classes=num_classes)
return model
predict code:
detections = model.predict(image_path, threshold=confidence_threshold)
bbox_list = detections.xyxy
score_list = detections.confidence
class_id_list = detections.class_id
shapes = []
for bbox, score, class_id in zip(bbox_list, score_list, class_id_list):
try:
class_name = meta_info[class_id]['class_name']
except KeyError:
print(class_id)

Additional

No response

Are you willing to submit a PR?

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

[wxqiankai]

I print vars(self.model) and results in RFDETR.predict(), find num_classes=36, However, the labels in the "results" section have values greater than 36. It seems that the "args" values of the model are not functioning properly. The model is still using the 90-class classification of COCO.

        with torch.inference_mode():
            print(vars(self.model))
            if self._is_optimized_for_inference:
                predictions = self.model.inference_model(batch_tensor.to(dtype=self._optimized_dtype))
            else:
                predictions = self.model.model(batch_tensor)
            if isinstance(predictions, tuple):
                predictions = {
                    "pred_logits": predictions[1],
                    "pred_boxes": predictions[0],
                    "pred_masks": predictions[2]
                }
            target_sizes = torch.tensor(orig_sizes, device=self.model.device)
            results = self.model.postprocess(predictions, target_sizes=target_sizes)
        print(results)

{'args': Namespace(num_classes=36,
'labels': tensor([ 4, 70, 26, 26, 26, 26, 26, 70, 35, 26, 70, 35, 70, 70, 35, 26, 35, 70,