运行visualize_cv2.py遇列表索引越界，求Mask-RCNN仅训练Person类方案

Let's break down the problem and fix it step by step—this error almost always pops up when your custom class list doesn't match the number of classes the model is configured to output. Here's how to get it working smoothly:

Why You're Seeing the "List index out of range" Error #

The pre-trained COCO Mask-RCNN model is built for 81 classes (including the background class "BG"). If you trimmed your class_names to just ['person'] (without including "BG"), or didn't adjust the model's output layers to match your new class count, the model will still spit out class indices that don't exist in your shortened list. For example, in COCO, "person" is class index 1—if your list only has ['person'], index 1 is out of bounds (since lists are 0-indexed).

Step 1: Correct Your Class List & Model Configuration #

First, define your class list to include the background class (required by Mask-RCNN's internal logic), and create a custom config that matches your new class count:

import cv2
import numpy as np
import os
import sys
import coco
import utils
import model as modellib
from config import Config

ROOT_DIR = os.getcwd()
MODEL_DIR = os.path.join(ROOT_DIR, "logs")
COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")

# 1. Define your class list (BG must be first)
class_names = ['BG', 'person']

# 2. Create a custom config for single-class training
class PersonConfig(Config):
    NAME = "person_detector"
    # Number of classes = BG + 1 target class (person)
    NUM_CLASSES = 1 + 1
    # Keep core configs aligned with COCO for consistency
    GPU_COUNT = 1
    IMAGES_PER_GPU = 2
    IMAGE_MIN_DIM = 800
    IMAGE_MAX_DIM = 1024

config = PersonConfig()
config.display()

Step 2: Load Pre-Trained Weights Correctly #

You can't load the full COCO weights directly—you need to exclude the final layers tied to the original 81 classes, so the model re-initializes these layers for your 2-class setup:

# Initialize model in training mode
model = modellib.MaskRCNN(mode="training", config=config, model_dir=MODEL_DIR)

# Load COCO weights, skipping class-specific layers
model.load_weights(COCO_MODEL_PATH, by_name=True, exclude=[
    "mrcnn_class_logits", "mrcnn_bbox_fc", 
    "mrcnn_bbox", "mrcnn_mask"
])

Step 3: Fix Inference/Visualization in visualize_cv2.py #

For visualization, use the same custom config in inference mode, and ensure the model loads weights that match your class count:

# Create inference config (adjust GPU settings if needed)
class InferenceConfig(PersonConfig):
    GPU_COUNT = 1
    IMAGES_PER_GPU = 1

inference_config = InferenceConfig()

# Initialize model in inference mode
model = modellib.MaskRCNN(mode="inference", config=inference_config, model_dir=MODEL_DIR)

# Load the latest trained weights (or adjusted pre-trained weights for testing)
model_path = model.find_last()
model.load_weights(model_path, by_name=True)

# Example prediction loop (no more index errors!)
image = cv2.imread("test_image.jpg")
results = model.detect([image], verbose=1)
r = results[0]

for i, class_id in enumerate(r['class_ids']):
    class_name = class_names[class_id]
    print(f"Detected: {class_name}")

Step 4: Prepare Your Training Data #

Make sure your dataset annotations are formatted correctly:

• All annotations should target the "person" class, with class ID 1 (since "BG" is ID 0)
• Follow the Mask-RCNN dataset format (use the COCO or Balloon dataset examples in the official repo as a template)

Final Quick Checks #

• Always start class_names with "BG"—this is non-negotiable for Mask-RCNN's logic
• Confirm every reference to class counts (config, model initialization, inference) uses the correct total (2 classes)

内容的提问来源于stack exchange，提问作者Dmitriy Kisil