How to Build Image Segmentation Tools with Mask R-CNN and PyTorch 2.3 for 2026

import os
import json
import torch
import numpy as np
from PIL import Image
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torchvision.transforms import functional as F
import logging
from pathlib import Path

# Configure logging for error tracking
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class COCOMaskDataset(Dataset):
    """Custom Dataset for loading COCO-format instance segmentation data with error handling."""
    def __init__(self, root_dir: str, annotation_file: str, transforms=None, min_box_size: int = 10):
        """
        Args:
            root_dir: Path to directory containing images
            annotation_file: Path to COCO-format annotation JSON
            transforms: Optional torchvision transforms for augmentation
            min_box_size: Minimum width/height for valid bounding boxes (filters tiny boxes)
        """
        self.root_dir = Path(root_dir)
        self.transforms = transforms
        self.min_box_size = min_box_size
        self.images = []
        self.annotations = {}
        self.img_id_to_file = {}

        # Load and validate COCO annotations
        try:
            with open(annotation_file, 'r') as f:
                coco_data = json.load(f)
        except (FileNotFoundError, json.JSONDecodeError) as e:
            logger.error(f'Failed to load annotation file {annotation_file}: {e}')
            raise

        # Map image IDs to file names
        for img in coco_data['images']:
            self.img_id_to_file[img['id']] = img['file_name']
        # Filter out images that don't exist on disk
        valid_img_ids = []
        for img_id, file_name in self.img_id_to_file.items():
            img_path = self.root_dir / file_name
            if img_path.exists():
                valid_img_ids.append(img_id)
            else:
                logger.warning(f'Image {file_name} not found on disk, skipping')

        # Group annotations by image ID
        for ann in coco_data['annotations']:
            img_id = ann['image_id']
            if img_id in valid_img_ids:
                if img_id not in self.annotations:
                    self.annotations[img_id] = []
                self.annotations[img_id].append(ann)

        self.images = list(self.annotations.keys())
        logger.info(f'Loaded {len(self.images)} valid images with annotations')

    def __len__(self) -> int:
        return len(self.images)

    def __getitem__(self, idx: int) -> tuple:
        """Returns (image, target) tuple with error handling for corrupt data."""
        img_id = self.images[idx]
        img_file = self.img_id_to_file[img_id]
        img_path = self.root_dir / img_file

        # Load image with error handling for corrupt files
        try:
            img = Image.open(img_path).convert('RGB')
        except (IOError, OSError) as e:
            logger.error(f'Corrupt image {img_path}: {e}, returning next sample')
            return self.__getitem__((idx + 1) % len(self))

        # Convert to tensor
        img = F.to_tensor(img)

        # Load annotations for this image
        anns = self.annotations[img_id]
        boxes = []
        labels = []
        masks = []
        valid_anns = 0

        for ann in anns:
            # Filter tiny bounding boxes
            x_min, y_min, w, h = ann['bbox']
            if w < self.min_box_size or h < self.min_box_size:
                continue
            # Filter invalid segmentations (at least 3 points for polygon)
            if len(ann['segmentation'][0]) < 6:
                continue
            boxes.append([x_min, y_min, x_min + w, y_min + h])
            labels.append(ann['category_id'])
            # Generate binary mask from COCO segmentation (simplified for example)
            mask = np.zeros((img.shape[1], img.shape[2]), dtype=np.uint8)
            masks.append(torch.from_numpy(mask))
            valid_anns += 1

        if valid_anns == 0:
            logger.warning(f'No valid annotations for image {img_id}, returning next sample')
            return self.__getitem__((idx + 1) % len(self))

        target = {
            'boxes': torch.tensor(boxes, dtype=torch.float32),
            'labels': torch.tensor(labels, dtype=torch.int64),
            'masks': torch.stack(masks) if masks else torch.zeros((0, img.shape[1], img.shape[2])),
            'image_id': torch.tensor([img_id])
        }

        if self.transforms:
            img, target = self.transforms(img, target)
        return img, target


def segmentation_collate(batch):
    """Custom collate function for variable-size segmentation targets."""
    images = [item[0] for item in batch]
    targets = [item[1] for item in batch]
    return images, targets


if __name__ == '__main__':
    # Example usage with COCO val2017
    dataset = COCOMaskDataset(
        root_dir='data/coco/val2017',
        annotation_file='data/coco/annotations/instances_val2017.json',
        min_box_size=10
    )
    dataloader = DataLoader(
        dataset,
        batch_size=4,
        shuffle=True,
        num_workers=2,
        collate_fn=segmentation_collate
    )
    for imgs, targets in dataloader:
        print(f'Batch image shapes: {[img.shape for img in imgs]}')
        print(f'Batch box counts: {[t["boxes"].shape[0] for t in targets]}')
        break

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor
from torch.cuda.amp import GradScaler, autocast
import logging
import os
from pathlib import Path

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def get_model(num_classes: int, pretrained: bool = True) -> nn.Module:
    """Initialize Mask R-CNN with custom head for num_classes (including background)."""
    # Load pretrained ResNet-50-FPN backbone
    model = maskrcnn_resnet50_fpn(pretrained=pretrained, progress=True)

    # Replace box predictor for custom number of classes
    in_features_box = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features_box, num_classes)

    # Replace mask predictor for custom number of classes
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, hidden_layer, num_classes)

    return model

def train_one_epoch(model, dataloader, optimizer, device, epoch, scaler):
    """Train for one epoch with AMP and gradient scaling."""
    model.train()
    total_loss = 0.0
    for batch_idx, (images, targets) in enumerate(dataloader):
        # Move data to device
        images = [img.to(device) for img in images]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        optimizer.zero_grad()

        # Forward pass with AMP
        with autocast():
            loss_dict = model(images, targets)
            losses = sum(loss for loss in loss_dict.values())

        # Check for NaN losses
        if torch.isnan(losses):
            logger.error(f'NaN loss encountered at batch {batch_idx}, skipping')
            continue

        # Backward pass with gradient scaling
        scaler.scale(losses).backward()
        scaler.step(optimizer)
        scaler.update()

        total_loss += losses.item()
        if batch_idx % 10 == 0:
            logger.info(f'Epoch {epoch}, Batch {batch_idx}, Loss: {losses.item():.4f}')

    return total_loss / len(dataloader)

if __name__ == '__main__':
    # Configuration
    NUM_CLASSES = 91  # 80 COCO classes + background
    BATCH_SIZE = 4
    LEARNING_RATE = 1e-4
    NUM_EPOCHS = 10
    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    CHECKPOINT_DIR = Path('checkpoints')
    CHECKPOINT_DIR.mkdir(exist_ok=True)

    logger.info(f'Training on device: {DEVICE}')

    # Initialize model, optimizer, scaler
    model = get_model(NUM_CLASSES)
    model.to(DEVICE)

    # Compile model with PyTorch 2.3 torch.compile for faster training
    try:
        model = torch.compile(model, mode='max-autotune')
        logger.info('Model compiled with max-autotune')
    except Exception as e:
        logger.warning(f'Failed to compile model: {e}, proceeding without compilation')

    optimizer = optim.AdamW(model.parameters(), lr=LEARNING_RATE)
    scaler = GradScaler()

    # Load dataset and dataloader (reuse COCOMaskDataset from Step 1)
    dataset = COCOMaskDataset(
        root_dir='data/coco/train2017',
        annotation_file='data/coco/annotations/instances_train2017.json'
    )
    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, collate_fn=segmentation_collate, num_workers=4)

    # Training loop
    for epoch in range(NUM_EPOCHS):
        avg_loss = train_one_epoch(model, dataloader, optimizer, DEVICE, epoch, scaler)
        logger.info(f'Epoch {epoch} average loss: {avg_loss:.4f}')

        # Save checkpoint
        checkpoint_path = CHECKPOINT_DIR / f'mask_rcnn_epoch_{epoch}.pth'
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': avg_loss
        }, checkpoint_path)
        logger.info(f'Saved checkpoint to {checkpoint_path}')

    logger.info('Training complete')

import torch
import torch.nn as nn
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision import transforms
from torch.ao.quantization import get_default_qconfig, prepare_qat, convert
import time
import numpy as np
from PIL import Image
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class SegmentationInference:
    """Production inference pipeline for Mask R-CNN with quantization support."""
    def __init__(self, checkpoint_path: str, num_classes: int = 91, device: str = 'cuda', quantized: bool = False):
        self.device = torch.device(device if torch.cuda.is_available() else 'cpu')
        self.num_classes = num_classes
        self.quantized = quantized

        # Load model
        self.model = self._load_model(checkpoint_path)
        self.model.to(self.device)
        self.model.eval()

        # Compile model for faster inference (PyTorch 2.3+)
        if not quantized:
            try:
                self.model = torch.compile(self.model, mode='reduce-overhead')
                logger.info('Model compiled for inference')
            except Exception as e:
                logger.warning(f'Failed to compile model: {e}')

        # Image preprocessing
        self.transform = transforms.Compose([
            transforms.ToTensor(),
        ])

    def _load_model(self, checkpoint_path: str) -> nn.Module:
        """Load model from checkpoint with error handling."""
        if not Path(checkpoint_path).exists():
            logger.error(f'Checkpoint {checkpoint_path} not found')
            raise FileNotFoundError(f'Checkpoint {checkpoint_path} not found')

        model = maskrcnn_resnet50_fpn(pretrained=False)
        in_features_box = model.roi_heads.box_predictor.cls_score.in_features
        model.roi_heads.box_predictor = FastRCNNPredictor(in_features_box, self.num_classes)
        in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
        model.roi_heads.mask_predictor = MaskRCNNPredictor(in_features_mask, 256, self.num_classes)

        checkpoint = torch.load(checkpoint_path, map_location=self.device)
        model.load_state_dict(checkpoint['model_state_dict'])
        logger.info(f'Loaded checkpoint from {checkpoint_path}')
        return model

    def quantize(self):
        """Apply post-training quantization to INT8 for edge deployment."""
        if self.quantized:
            logger.warning('Model already quantized')
            return

        # Prepare model for quantization
        self.model.qconfig = get_default_qconfig('x86')  # Use 'fbgemm' for x86, 'qnnpack' for ARM
        prepare_qat(self.model, inplace=True)

        # Calibrate with sample data (simplified; use 100+ real samples in production)
        calib_img = torch.randn(1, 3, 640, 640).to(self.device)
        self.model(calib_img)

        # Convert to quantized model
        convert(self.model, inplace=True)
        self.quantized = True
        logger.info('Model quantized to INT8')

    def infer(self, image_path: str) -> dict:
        """Run inference on a single image."""
        if not Path(image_path).exists():
            raise FileNotFoundError(f'Image {image_path} not found')

        # Load and preprocess image
        img = Image.open(image_path).convert('RGB')
        img_tensor = self.transform(img).unsqueeze(0).to(self.device)

        # Run inference
        with torch.no_grad():
            start = time.time()
            predictions = self.model(img_tensor)[0]
            latency = (time.time() - start) * 1000  # ms

        # Postprocess predictions
        masks = predictions['masks'].cpu().numpy()
        boxes = predictions['boxes'].cpu().numpy()
        labels = predictions['labels'].cpu().numpy()
        scores = predictions['scores'].cpu().numpy()

        return {
            'masks': masks,
            'boxes': boxes,
            'labels': labels,
            'scores': scores,
            'latency_ms': latency
        }

    def benchmark(self, image_dir: str, num_runs: int = 100) -> dict:
        """Benchmark inference latency over multiple runs."""
        image_paths = list(Path(image_dir).glob('*.jpg'))[:10]  # Use 10 images for benchmarking
        latencies = []

        for _ in range(num_runs):
            for img_path in image_paths:
                try:
                    result = self.infer(str(img_path))
                    latencies.append(result['latency_ms'])
                except Exception as e:
                    logger.error(f'Inference failed on {img_path}: {e}')

        return {
            'p50_latency_ms': np.percentile(latencies, 50),
            'p99_latency_ms': np.percentile(latencies, 99),
            'avg_latency_ms': np.mean(latencies),
            'throughput_inf_s': 1000 / np.mean(latencies)
        }

if __name__ == '__main__':
    # Example usage
    inference = SegmentationInference(
        checkpoint_path='checkpoints/mask_rcnn_epoch_9.pth',
        num_classes=91,
        device='cuda'
    )

    # Benchmark unquantized model
    print('Unquantized benchmark:')
    unquant_bench = inference.benchmark('data/coco/val2017', num_runs=10)
    print(f'p50 latency: {unquant_bench["p50_latency_ms"]:.2f}ms')
    print(f'p99 latency: {unquant_bench["p99_latency_ms"]:.2f}ms')
    print(f'Throughput: {unquant_bench["throughput_inf_s"]:.2f} inf/s')

    # Quantize and benchmark
    inference.quantize()
    print('\nQuantized benchmark:')
    quant_bench = inference.benchmark('data/coco/val2017', num_runs=10)
    print(f'p50 latency: {quant_bench["p50_latency_ms"]:.2f}ms')
    print(f'p99 latency: {quant_bench["p99_latency_ms"]:.2f}ms')
    print(f'Throughput: {quant_bench["throughput_inf_s"]:.2f} inf/s')

try:
    model = torch.compile(model, mode='max-autotune')
except Exception as e:
    logger.warning(f'Compilation failed: {e}, using reduce-overhead')
    model = torch.compile(model, mode='reduce-overhead')

from pycocotools.coco import COCO
def validate_coco_annotations(annotation_file):
    coco = COCO(annotation_file)
    img_ids = coco.getImgIds()
    for img_id in img_ids:
        ann_ids = coco.getAnnIds(imgIds=img_id)
        anns = coco.loadAnns(ann_ids)
        for ann in anns:
            if ann['bbox'][2] <=0 or ann['bbox'][3] <=0:
                print(f'Invalid bbox for ann {ann["id"]}')

from torch.ao.quantization import get_default_qat_qconfig, prepare_qat, convert
model.qconfig = get_default_qat_qconfig('x86')
prepare_qat(model, inplace=True)
# Train for 2-3 more epochs with QAT
for epoch in range(2):
    train_one_epoch(model, dataloader, optimizer, DEVICE, epoch, scaler)
convert(model, inplace=True)

mask-rcnn-2026/
├── data/
│   ├── coco/
│   ├── custom/
│   └── preprocess.py
├── models/
│   ├── mask_rcnn.py
│   ├── backbone.py
│   └── checkpoint.pth
├── training/
│   ├── train.py
│   ├── config.yaml
│   └── loss.py
├── inference/
│   ├── infer.py
│   ├── quantize.py
│   └── benchmark.py
├── api/
│   ├── main.py
│   └── requirements.txt
├── tests/
│   ├── test_data.py
│   └── test_model.py
├── requirements.txt
└── README.md