AssertionError during evaluation: Predicted class ID out of range in Detectron2 RetinaNet fine-tuning

Question

I am fine-tuning a Detectron2 RetinaNet model with a ResNet-50 backbone using my custom dataset, which contains 3 classes (excluding the background).

The training process works fine, and the model produces reasonable results during inference. However, when I evaluate the model, I encounter the following error:

AssertionError: A prediction has class=10, but the dataset only has 4 classes and predicted class id should be in [0, 3].

Here is my full code:

!pip install pyyaml==5.1

import torch
TORCH_VERSION = ".".join(torch.__version__.split(".")[:2])
CUDA_VERSION = torch.__version__.split("+")[-1]
print("torch: ", TORCH_VERSION, "; cuda: ", CUDA_VERSION)
# Install detectron2 that matches the above pytorch version
# See https://detectron2.readthedocs.io/tutorials/install.html for instructions
# !pip install detectron2 -f https://dl.fbaipublicfiles.com/detectron2/wheels/TORCH_VERSION/index.html
# If there is not yet a detectron2 release that matches the given torch + CUDA version, you need to install a different pytorch.
!python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'
# exit(0)  # After installation, you may need to "restart runtime" in Colab. This line can also restart runtime

# Some basic setup:
# Setup detectron2 logger
import detectron2
from detectron2.utils.logger import setup_logger
setup_logger()

# import some common libraries
import numpy as np
import os, json, cv2, random
from google.colab.patches import cv2_imshow

# import some common detectron2 utilities
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2.utils.visualizer import Visualizer
from detectron2.data import MetadataCatalog, DatasetCatalog

#visualize training data
my_dataset_train_metadata = MetadataCatalog.get("my_dataset_train")
dataset_dicts = DatasetCatalog.get("my_dataset_train")

import random
from detectron2.utils.visualizer import Visualizer

for d in random.sample(dataset_dicts, 3):
    img = cv2.imread(d["file_name"])
    visualizer = Visualizer(img[:, :, ::-1], metadata=my_dataset_train_metadata, scale=0.5)
    vis = visualizer.draw_dataset_dict(d)
    cv2_imshow(vis.get_image()[:, :, ::-1])


from detectron2.engine import DefaultTrainer

cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-Detection/retinanet_R_50_FPN_1x.yaml"))
cfg.DATASETS.TRAIN = ("my_dataset_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 2
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-Detection/retinanet_R_50_FPN_1x.yaml")  # Let training initialize from model zoo
cfg.SOLVER.IMS_PER_BATCH = 8
cfg.SOLVER.BASE_LR = 0.0025  # pick a good LR
cfg.SOLVER.MAX_ITER = 1632   # 300 iterations seems good enough for this toy dataset; you will need to train longer for a practical dataset
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 128  # faster, and good enough for this toy dataset (default: 512)
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 3  # (see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets)
# NOTE: this config means the number of classes, but a few popular unofficial tutorials incorrect uses num_classes+1 here.

os.makedirs(cfg.OUTPUT_DIR, exist_ok=True)
trainer = DefaultTrainer(cfg)
trainer.resume_or_load(resume=True)
trainer.train()

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.6   # set the testing threshold for this model
cfg.DATASETS.TEST = ("my_dataset_test", )
predictor = DefaultPredictor(cfg)

test_metadata = MetadataCatalog.get("my_dataset_test")

Inferencing Block:

from detectron2.utils.visualizer import Visualizer, ColorMode
import cv2
import glob
from google.colab.patches import cv2_imshow  # For Colab

# Adjust the SCORE_THRESH_TEST to filter low-confidence predictions
cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "/content/drive/MyDrive/output/model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.50  # Confidence threshold for predictions
cfg.DATASETS.TEST = ("my_dataset_test",)
predictor = DefaultPredictor(cfg)
test_metadata = MetadataCatalog.get("my_dataset_test")

def non_max_suppression(instances, iou_threshold=0.5):
    """
    Apply Non-Maximum Suppression (NMS) to filter overlapping boxes.
    """
    import torch
    boxes = instances.pred_boxes.tensor.cpu()
    scores = instances.scores.cpu()
    keep = torch.ops.torchvision.nms(boxes, scores, iou_threshold)
    return instances[keep]

for imageName in glob.glob('/content/Capstone_Team_Nugget_Final-4/test/Potato_Late_blight_raw-597-_JPG.rf.84dda4b4e1c2fa556f6a38b0e5b70b19.jpg'):
    im = cv2.imread(imageName)
    outputs = predictor(im)

    # Apply NMS to remove overlapping boxes
    instances = outputs["instances"].to("cpu")

    # Filter low-confidence predictions
    high_conf_instances = instances[instances.scores > 0.50]
    final_instances = non_max_suppression(high_conf_instances, iou_threshold=0.5)

    # Visualization in RGB
    v = Visualizer(
        im[:, :, ::-1],
        metadata=test_metadata,
        scale=0.8,
        instance_mode=ColorMode.IMAGE  # Use IMAGE for RGB visualization
    )
    out = v.draw_instance_predictions(final_instances)
    cv2_imshow(out.get_image()[:, :, ::-1])

Evaluation Block:

from detectron2.evaluation import COCOEvaluator, inference_on_dataset
from detectron2.data import build_detection_test_loader
evaluator = COCOEvaluator("my_dataset_test", output_dir="/content/drive/MyDrive/output")
val_loader = build_detection_test_loader(cfg, "my_dataset_test")
print(val_loader)
print(inference_on_dataset(predictor.model, val_loader, evaluator))

StackTrace:

WARNING [01/25 22:14:30 d2.data.datasets.coco]: 
Category ids in annotations are not in [1, #categories]! We'll apply a mapping for you.

[01/25 22:14:30 d2.data.datasets.coco]: Loaded 1091 images in COCO format from /content/Capstone_Team_Nugget_Final-4/valid/_annotations.coco.json
[01/25 22:14:30 d2.data.dataset_mapper]: [DatasetMapper] Augmentations used in inference: [ResizeShortestEdge(short_edge_length=(800, 800), max_size=1333, sample_style='choice')]
[01/25 22:14:30 d2.data.common]: Serializing the dataset using: <class 'detectron2.data.common._TorchSerializedList'>
[01/25 22:14:30 d2.data.common]: Serializing 1091 elements to byte tensors and concatenating them all ...
[01/25 22:14:30 d2.data.common]: Serialized dataset takes 0.40 MiB
<torch.utils.data.dataloader.DataLoader object at 0x7ea5f1b09390>
[01/25 22:14:30 d2.evaluation.evaluator]: Start inference on 1091 batches
[01/25 22:14:32 d2.evaluation.evaluator]: Inference done 11/1091. Dataloading: 0.0013 s/iter. Inference: 0.0751 s/iter. Eval: 0.0003 s/iter. Total: 0.0767 s/iter. ETA=0:01:22
[01/25 22:14:37 d2.evaluation.evaluator]: Inference done 76/1091. Dataloading: 0.0018 s/iter. Inference: 0.0751 s/iter. Eval: 0.0003 s/iter. Total: 0.0772 s/iter. ETA=0:01:18
[01/25 22:14:42 d2.evaluation.evaluator]: Inference done 141/1091. Dataloading: 0.0020 s/iter. Inference: 0.0753 s/iter. Eval: 0.0003 s/iter. Total: 0.0777 s/iter. ETA=0:01:13
[01/25 22:14:47 d2.evaluation.evaluator]: Inference done 201/1091. Dataloading: 0.0032 s/iter. Inference: 0.0758 s/iter. Eval: 0.0003 s/iter. Total: 0.0794 s/iter. ETA=0:01:10
[01/25 22:14:52 d2.evaluation.evaluator]: Inference done 266/1091. Dataloading: 0.0029 s/iter. Inference: 0.0758 s/iter. Eval: 0.0003 s/iter. Total: 0.0790 s/iter. ETA=0:01:05
[01/25 22:14:57 d2.evaluation.evaluator]: Inference done 330/1091. Dataloading: 0.0028 s/iter. Inference: 0.0759 s/iter. Eval: 0.0003 s/iter. Total: 0.0790 s/iter. ETA=0:01:00
[01/25 22:15:02 d2.evaluation.evaluator]: Inference done 390/1091. Dataloading: 0.0033 s/iter. Inference: 0.0762 s/iter. Eval: 0.0003 s/iter. Total: 0.0798 s/iter. ETA=0:00:55
[01/25 22:15:07 d2.evaluation.evaluator]: Inference done 454/1091. Dataloading: 0.0031 s/iter. Inference: 0.0762 s/iter. Eval: 0.0003 s/iter. Total: 0.0796 s/iter. ETA=0:00:50
[01/25 22:15:12 d2.evaluation.evaluator]: Inference done 518/1091. Dataloading: 0.0029 s/iter. Inference: 0.0763 s/iter. Eval: 0.0003 s/iter. Total: 0.0796 s/iter. ETA=0:00:45
[01/25 22:15:17 d2.evaluation.evaluator]: Inference done 578/1091. Dataloading: 0.0032 s/iter. Inference: 0.0765 s/iter. Eval: 0.0003 s/iter. Total: 0.0801 s/iter. ETA=0:00:41
[01/25 22:15:22 d2.evaluation.evaluator]: Inference done 642/1091. Dataloading: 0.0030 s/iter. Inference: 0.0766 s/iter. Eval: 0.0003 s/iter. Total: 0.0800 s/iter. ETA=0:00:35
[01/25 22:15:27 d2.evaluation.evaluator]: Inference done 705/1091. Dataloading: 0.0030 s/iter. Inference: 0.0767 s/iter. Eval: 0.0003 s/iter. Total: 0.0800 s/iter. ETA=0:00:30
[01/25 22:15:32 d2.evaluation.evaluator]: Inference done 764/1091. Dataloading: 0.0033 s/iter. Inference: 0.0768 s/iter. Eval: 0.0003 s/iter. Total: 0.0805 s/iter. ETA=0:00:26
[01/25 22:15:37 d2.evaluation.evaluator]: Inference done 827/1091. Dataloading: 0.0032 s/iter. Inference: 0.0769 s/iter. Eval: 0.0003 s/iter. Total: 0.0804 s/iter. ETA=0:00:21
[01/25 22:15:42 d2.evaluation.evaluator]: Inference done 877/1091. Dataloading: 0.0034 s/iter. Inference: 0.0777 s/iter. Eval: 0.0003 s/iter. Total: 0.0815 s/iter. ETA=0:00:17
[01/25 22:15:47 d2.evaluation.evaluator]: Inference done 933/1091. Dataloading: 0.0037 s/iter. Inference: 0.0779 s/iter. Eval: 0.0003 s/iter. Total: 0.0820 s/iter. ETA=0:00:12
[01/25 22:15:52 d2.evaluation.evaluator]: Inference done 994/1091. Dataloading: 0.0037 s/iter. Inference: 0.0779 s/iter. Eval: 0.0003 s/iter. Total: 0.0821 s/iter. ETA=0:00:07
[01/25 22:15:57 d2.evaluation.evaluator]: Inference done 1057/1091. Dataloading: 0.0036 s/iter. Inference: 0.0779 s/iter. Eval: 0.0003 s/iter. Total: 0.0819 s/iter. ETA=0:00:02
[01/25 22:16:00 d2.evaluation.evaluator]: Total inference time: 0:01:29.050782 (0.081999 s / iter per device, on 1 devices)
[01/25 22:16:00 d2.evaluation.evaluator]: Total inference pure compute time: 0:01:24 (0.077938 s / iter per device, on 1 devices)
[01/25 22:16:00 d2.evaluation.coco_evaluation]: Preparing results for COCO format ...
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-92-e539605f50ff> in <cell line: 0>()
      4 val_loader = build_detection_test_loader(cfg, "my_dataset_test")
      5 print(val_loader)
----> 6 print(inference_on_dataset(predictor.model, val_loader, evaluator))

2 frames
/usr/local/lib/python3.11/dist-packages/detectron2/evaluation/evaluator.py in inference_on_dataset(model, data_loader, evaluator, callbacks)
    213     )
    214 
--> 215     results = evaluator.evaluate()
    216     # An evaluator may return None when not in main process.
    217     # Replace it by an empty dict instead to make it easier for downstream code to handle

/usr/local/lib/python3.11/dist-packages/detectron2/evaluation/coco_evaluation.py in evaluate(self, img_ids)
    204             self._eval_box_proposals(predictions)
    205         if "instances" in predictions[0]:
--> 206             self._eval_predictions(predictions, img_ids=img_ids)
    207         # Copy so the caller can do whatever with results
    208         return copy.deepcopy(self._results)

/usr/local/lib/python3.11/dist-packages/detectron2/evaluation/coco_evaluation.py in _eval_predictions(self, predictions, img_ids)
    238             for result in coco_results:
    239                 category_id = result["category_id"]
--> 240                 assert category_id < num_classes, (
    241                     f"A prediction has class={category_id}, "
    242                     f"but the dataset only has {num_classes} classes and "

AssertionError: A prediction has class=10, but the dataset only has 4 classes and predicted class id should be in [0, 3].