Mask R-CNN自定义带掩码数据集训练困惑，寻求技术指导

Hey there! I get where you're coming from—adapting Mask R-CNN implementations like Detectron from the standard COCO dataset to your own custom dataset (with per-image masks) can feel overwhelming at first, but it’s totally manageable with clear steps. Let’s break down the key areas to address your confusion:

1. Align Your Dataset with Framework Expectations #

Most Mask R-CNN frameworks (including Detectron/Detectron2) are built around COCO’s annotation structure, so you’ll need to map your custom data to a compatible format:

• Core Annotation Components: Each entry needs to include:
  • Image metadata: Unique ID, file path, width/height
  • Object annotations: Bounding box coordinates (bbox), category ID, and segmentation mask (either as RLE encoding or polygon coordinates—avoid raw pixel masks for efficiency)
• Detectron2 Shortcut: You don’t have to build a full COCO JSON. Instead, write a custom loader function that returns dictionaries containing the required fields, then register it with DatasetCatalog and MetadataCatalog to make the framework recognize your dataset.

2. Convert Your Mask Format #

Custom masks are often stored as PNGs (with pixel values representing class IDs) or exported as polygons from tools like LabelMe. Here’s how to convert them for training:

• PNG Masks to RLE: Use pycocotools.mask.encode() to convert binary masks (one per object) into RLE format—this is the space-efficient standard used by COCO.
• Polygons Directly: If your annotation tool outputs polygon coordinates (e.g., [(x1,y1), (x2,y2), ...]), you can pass these directly as the segmentation field in your annotations; frameworks like Detectron2 will handle the rest.
• Critical Check: Always verify that masks are perfectly aligned with their corresponding images (same dimensions, no cropping/offset issues)—misalignment will tank model performance.

3. Framework-Specific Setup (Detectron2 Example) #

Let’s use Detectron2 since it’s one of the most popular Mask R-CNN implementations:

• Register Your Dataset:

  from detectron2.data import DatasetCatalog, MetadataCatalog
  
  def load_custom_data():
      # Write logic to load your images and annotations here
      # Return a list of dicts, each with:
      # "file_name", "image_id", "height", "width", "annotations" (list of dicts with "bbox", "category_id", "segmentation")
      return my_dataset_list
  
  DatasetCatalog.register("custom_train", load_custom_data)
  MetadataCatalog.get("custom_train").set(thing_classes=["cat", "dog", "bird"]) # Replace with your classes
  

• Adjust Configuration:
  • Update DATASETS.TRAIN and DATASETS.TEST to your registered dataset names
  • Set MODEL.ROI_HEADS.NUM_CLASSES to the number of object classes in your dataset (note: background is handled automatically in Detectron2, so don’t add 1 here)
  • Tweak data augmentations (e.g., INPUT.AUGMENTATIONS) to match your dataset’s needs (random flips, scaling help with generalization)

4. Optimize Training for Your Dataset #

• Use Pre-Trained Weights: Start with COCO pre-trained weights (e.g., mask_rcnn_R_50_FPN_3x.yaml in Detectron2) instead of training from scratch—this drastically speeds up convergence and avoids overfitting on small datasets.
• Adjust Learning Rate: Learning rate scales with batch size. If your batch size is 1/8 of COCO’s typical 16, set your initial learning rate to ~0.0025 (down from COCO’s 0.02).
• Add Validation: Include a validation split to track metrics like mAP (object detection) and mask IoU (segmentation) during training—this helps you catch overfitting early.

5. Troubleshoot Common Pitfalls #

• Mask Alignment Errors: If predictions are misaligned, double-check that image/mask dimensions match and that bounding boxes correctly enclose their masks.
• Category ID Issues: Ensure your category IDs are continuous (e.g., 0,1,2 for 3 classes) and match the order in thing_classes.
• Memory Bloat: Raw pixel masks eat up memory—stick to RLE or polygon formats to keep training feasible.

内容的提问来源于stack exchange，提问作者Vajira Prabuddhaka