Error when trying to visualize output with custom dataset

[ekorudiawan]

I am trying to train custom dataset with RetinaNet model but got an error when trying to visualize output from the predictor. I have trying an example with ballon dataset and mask RCNN without any problem. The problem look like fail to get class name when calling function draw_instance_predictions

Here is the code for loading dataset

def get_robocup_dicts(img_dir):
    json_file = os.path.join(img_dir, "train.json")
    with open(json_file) as f:
        imgs_anns = json.load(f)

    dataset_dicts = []
    for idx, v in enumerate(imgs_anns['images']):
        record = {}
        filename = os.path.join(img_dir, v["file_name"])

        record["file_name"] = filename
        record["id"] = v["id"]
        record["height"] = v["height"]
        record["width"] = v["width"]

        annos = imgs_anns["annotations"]
        objs = []
        for anno in annos:
            if anno['image_id'] == v["id"]:
                obj = {
                    "bbox": anno['bbox'],
                    "bbox_mode": BoxMode.XYWH_ABS,
                    "segmentation": anno['segmentation'],
                    "category_id": anno['category_id'],
                    "iscrowd": 0
                }
                objs.append(obj)
        record["annotations"] = objs
        dataset_dicts.append(record)
        print(idx)
    return dataset_dicts

TRAIN_PATH = "../RoboCup-Dataset/"
for d in ["train"]:
    DatasetCatalog.register("robocup_" + d, lambda d=d: get_robocup_dicts(TRAIN_PATH + d))
    MetadataCatalog.get("robocup_" + d).set(thing_classes=["ball", "field", "left_goal", "right_goal"])
robocup_metadata = MetadataCatalog.get("robocup_train")
from detectron2.engine import DefaultTrainer
from detectron2.config import get_cfg

cfg = get_cfg()
cfg.merge_from_file("../../detectron2/configs/COCO-Detection/retinanet_R_50_FPN_1x.yaml")
cfg.DATASETS.TRAIN = ("robocup_train",)
cfg.DATASETS.TEST = ()
cfg.DATALOADER.NUM_WORKERS = 0
cfg.MODEL.WEIGHTS = "detectron2://COCO-Detection/retinanet_R_50_FPN_1x/137593951/model_final_b796dc.pkl"  
cfg.SOLVER.IMS_PER_BATCH = 2
cfg.SOLVER.BASE_LR = 0.00025
cfg.SOLVER.MAX_ITER = 300   
cfg.MODEL.ROI_HEADS.BATCH_SIZE_PER_IMAGE = 2   
cfg.MODEL.ROI_HEADS.NUM_CLASSES = 4  

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

Until this step, there is no problem, then test to visualize the output

cfg.MODEL.WEIGHTS = os.path.join(cfg.OUTPUT_DIR, "model_final.pth")
cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5   # set the testing threshold for this model
cfg.DATASETS.TEST = ("robocup_train")
predictor = DefaultPredictor(cfg)
from detectron2.utils.visualizer import ColorMode
dataset_dicts = get_robocup_dicts(TRAIN_PATH+"train")
for d in random.sample(dataset_dicts, 3):    
    im = cv2.imread(d["file_name"])
    outputs = predictor(im)
    v = Visualizer(im[:, :, ::-1],
                   metadata=robocup_metadata, 
                   scale=0.8
    )
    v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
    plt.imshow(v.get_image()[:, :, ::-1])
    plt.show()

Error in line

IndexError                                Traceback (most recent call last)
<ipython-input-16-a3ec4d15c6ce> in <module>
      8                    scale=0.8
      9     )
---> 10     v = v.draw_instance_predictions(outputs["instances"].to("cpu"))
     11     plt.imshow(v.get_image()[:, :, ::-1])
     12     plt.show()

/home/Detectron2-RoboCup/detectron2/detectron2/utils/visualizer.py in draw_instance_predictions(self, predictions)
    335         scores = predictions.scores if predictions.has("scores") else None
    336         classes = predictions.pred_classes if predictions.has("pred_classes") else None
--> 337         labels = _create_text_labels(classes, scores, self.metadata.get("thing_classes", None))
    338         keypoints = predictions.pred_keypoints if predictions.has("pred_keypoints") else None
    339 

/home/Detectron2-RoboCup/detectron2/detectron2/utils/visualizer.py in _create_text_labels(classes, scores, class_names)
    199     labels = None
    200     if classes is not None and class_names is not None and len(class_names) > 1:
--> 201         labels = [class_names[i] for i in classes]
    202     if scores is not None:
    203         if labels is None:

/home/Detectron2-RoboCup/detectron2/detectron2/utils/visualizer.py in <listcomp>(.0)
    199     labels = None
    200     if classes is not None and class_names is not None and len(class_names) > 1:
--> 201         labels = [class_names[i] for i in classes]
    202     if scores is not None:
    203         if labels is None:

IndexError: list index out of range

Test get classes from metadata

robocup_metadata.get("thing_classes")

got

['ball', 'field', 'left_goal', 'right_goal']

Environment

------------------------  -------------------------------------------------------------------
sys.platform              linux
Python                    3.6.9 |Anaconda, Inc.| (default, Jul 30 2019, 19:07:31) [GCC 7.3.0]
Numpy                     1.17.2
Detectron2 Compiler       GCC 7.4
Detectron2 CUDA Compiler  10.1
DETECTRON2_ENV_MODULE     <not set>
PyTorch                   1.3.0a0+24ae9b5
PyTorch Debug Build       False
torchvision               0.5.0a0
CUDA available            True
GPU 0                     GeForce GTX 1070
CUDA_HOME                 /usr/local/cuda
NVCC                      Cuda compilation tools, release 10.1, V10.1.243
TORCH_CUDA_ARCH_LIST      Kepler;Kepler+Tesla;Maxwell;Maxwell+Tegra;Pascal;Volta;Turing
Pillow                    6.2.1
cv2                       3.4.1
------------------------  -------------------------------------------------------------------
PyTorch built with:
  - GCC 7.4
  - Intel(R) Math Kernel Library Version 2019.0.1 Product Build 20180928 for Intel(R) 64 architecture applications
  - Intel(R) MKL-DNN v0.20.3 (Git Hash N/A)
  - OpenMP 201511 (a.k.a. OpenMP 4.5)
  - NNPACK is enabled
  - CUDA Runtime 10.1
  - NVCC architecture flags: -gencode;arch=compute_52,code=sm_52;-gencode;arch=compute_60,code=sm_60;-gencode;arch=compute_61,code=sm_61;-gencode;arch=compute_70,code=sm_70;-gencode;arch=compute_75,code=sm_75;-gencode;arch=compute_75,code=compute_75
  - CuDNN 7.6.4
  - Magma 2.1.0
  - Build settings: BLAS=MKL, BUILD_NAMEDTENSOR=OFF, BUILD_TYPE=Release, CXX_FLAGS= -fvisibility-inlines-hidden -fopenmp -DUSE_FBGEMM -DUSE_QNNPACK -O2 -fPIC -Wno-narrowing -Wall -Wextra -Wno-missing-field-initializers -Wno-type-limits -Wno-array-bounds -Wno-unknown-pragmas -Wno-sign-compare -Wno-unused-parameter -Wno-unused-variable -Wno-unused-function -Wno-unused-result -Wno-strict-overflow -Wno-strict-aliasing -Wno-error=deprecated-declarations -Wno-stringop-overflow -Wno-error=pedantic -Wno-error=redundant-decls -Wno-error=old-style-cast -fdiagnostics-color=always -faligned-new -Wno-unused-but-set-variable -Wno-maybe-uninitialized -fno-math-errno -fno-trapping-math -Wno-stringop-overflow, FORCE_FALLBACK_CUDA_MPI=1, PERF_WITH_AVX=1, PERF_WITH_AVX2=1, PERF_WITH_AVX512=1, USE_CUDA=True, USE_EXCEPTION_PTR=1, USE_GFLAGS=OFF, USE_GLOG=OFF, USE_MKL=ON, USE_MKLDNN=ON, USE_MPI=ON, USE_NCCL=ON, USE_NNPACK=ON, USE_OPENMP=ON, USE_STATIC_DISPATCH=OFF, 

[davidbunk]

I had a similar problem, it turned out that the IDs of my categories started at 1 and not as expected at 0. So if you have 4 classes, your category ID range might go 1-4, while the range of the name array with 4 entries goes from 0-3.

If this is the case, you can map your category IDs to the proper range by following instructions at https://detectron2.readthedocs.io/tutorials/datasets.html for setting thing_dataset_id_to_contiguous_id in the metadata of your dataset.

[ppwwyyxx]

You have not updated the config properly for RetinaNet. Please see https://detectron2.readthedocs.io/tutorials/datasets.html#update-the-config-for-new-datasets.