Instance Segmentation

1import torch
2import torchvision
3from torchvision.models.detection import maskrcnn_resnet50_fpn
4import numpy as np
5import cv2
6
7class SegmentationTypes:
8    """Demonstration of different segmentation approaches."""
9
10    def __init__(self):
11        # Load pretrained models
12        self.instance_model = maskrcnn_resnet50_fpn(pretrained=True)
13        self.instance_model.eval()
14
15        # COCO class names for instance segmentation
16        self.coco_names = [
17            '__background__', 'person', 'bicycle', 'car', 'motorcycle',
18            'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
19            'fire hydrant', 'N/A', 'stop sign', 'parking meter', 'bench',
20            'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant',
21            'bear', 'zebra', 'giraffe', 'N/A', 'backpack', 'umbrella'
22        ]
23
24    def semantic_segmentation(self, image: np.ndarray) -> np.ndarray:
25        """
26        Semantic segmentation - each pixel gets a class label.
27        All cars are labeled the same, no distinction between instances.
28        """
29        # Using DeepLabV3 for semantic segmentation
30        model = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True)
31        model.eval()
32
33        # Preprocess
34        transform = torchvision.transforms.Compose([
35            torchvision.transforms.ToTensor(),
36            torchvision.transforms.Normalize(
37                mean=[0.485, 0.456, 0.406],
38                std=[0.229, 0.224, 0.225]
39            )
40        ])
41
42        input_tensor = transform(image).unsqueeze(0)
43
44        with torch.no_grad():
45            output = model(input_tensor)['out'][0]
46
47        # Get class predictions for each pixel
48        semantic_mask = output.argmax(0).numpy()
49        return semantic_mask
50
51    def instance_segmentation(self, image: np.ndarray,
52                             confidence_threshold: float = 0.5):
53        """
54        Instance segmentation - each object gets unique ID and mask.
55        Two cars are distinguished as separate instances.
56        """
57        transform = torchvision.transforms.ToTensor()
58        input_tensor = transform(image).unsqueeze(0)
59
60        with torch.no_grad():
61            predictions = self.instance_model(input_tensor)[0]
62
63        # Filter by confidence
64        high_conf_idx = predictions['scores'] > confidence_threshold
65
66        masks = predictions['masks'][high_conf_idx].squeeze(1).numpy()
67        labels = predictions['labels'][high_conf_idx].numpy()
68        scores = predictions['scores'][high_conf_idx].numpy()
69        boxes = predictions['boxes'][high_conf_idx].numpy()
70
71        return {
72            'masks': masks,
73            'labels': [self.coco_names[l] for l in labels],
74            'scores': scores,
75            'boxes': boxes
76        }
77
78    def panoptic_segmentation(self, semantic_mask: np.ndarray,
79                             instance_results: dict) -> np.ndarray:
80        """
81        Panoptic segmentation combines semantic and instance.
82        - 'Stuff' classes (sky, road) from semantic segmentation
83        - 'Things' classes (cars, people) from instance segmentation
84        """
85        h, w = semantic_mask.shape
86        panoptic_mask = np.zeros((h, w), dtype=np.int32)
87
88        # First, apply semantic segmentation for 'stuff' classes
89        stuff_classes = [0, 14, 15, 16, 17, 18, 19, 20]  # background, ground, etc.
90        for stuff_id in stuff_classes:
91            panoptic_mask[semantic_mask == stuff_id] = stuff_id
92
93        # Then overlay instance masks for 'things'
94        instance_id_offset = 1000
95        for idx, mask in enumerate(instance_results['masks']):
96            binary_mask = mask > 0.5
97            instance_id = instance_id_offset + idx
98            panoptic_mask[binary_mask] = instance_id
99
100        return panoptic_mask
101
102# Usage example
103segmenter = SegmentationTypes()
104image = cv2.imread('street_scene.jpg')
105image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
106
107# Get different segmentation types
108semantic = segmenter.semantic_segmentation(image_rgb)
109instance = segmenter.instance_segmentation(image_rgb)
110print(f"Found {len(instance['labels'])} instances: {instance['labels']}")

1import torch
2import torch.nn as nn
3import torchvision
4from torchvision.models.detection import MaskRCNN
5from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
6from torchvision.models.detection.rpn import AnchorGenerator
7
8class CustomMaskRCNN:
9    """
10    Custom Mask R-CNN with configurable backbone and heads.
11    """
12
13    def __init__(self, num_classes: int, backbone_name: str = 'resnet50'):
14        self.num_classes = num_classes
15        self.backbone_name = backbone_name
16        self.model = self._build_model()
17
18    def _build_model(self) -> MaskRCNN:
19        """Build Mask R-CNN with custom configuration."""
20
21        # Build backbone with FPN
22        backbone = resnet_fpn_backbone(
23            backbone_name=self.backbone_name,
24            pretrained=True,
25            trainable_layers=3  # Freeze early layers
26        )
27
28        # Custom anchor generator for different object scales
29        anchor_generator = AnchorGenerator(
30            sizes=((32,), (64,), (128,), (256,), (512,)),
31            aspect_ratios=((0.5, 1.0, 2.0),) * 5
32        )
33
34        # RoI pooling output size
35        roi_pooler = torchvision.ops.MultiScaleRoIAlign(
36            featmap_names=['0', '1', '2', '3'],
37            output_size=7,
38            sampling_ratio=2
39        )
40
41        # Mask RoI pooling (higher resolution for masks)
42        mask_roi_pooler = torchvision.ops.MultiScaleRoIAlign(
43            featmap_names=['0', '1', '2', '3'],
44            output_size=14,
45            sampling_ratio=2
46        )
47
48        # Build full model
49        model = MaskRCNN(
50            backbone,
51            num_classes=self.num_classes,
52            rpn_anchor_generator=anchor_generator,
53            box_roi_pool=roi_pooler,
54            mask_roi_pool=mask_roi_pooler,
55            # Detection parameters
56            rpn_pre_nms_top_n_train=2000,
57            rpn_pre_nms_top_n_test=1000,
58            rpn_post_nms_top_n_train=2000,
59            rpn_post_nms_top_n_test=1000,
60            rpn_nms_thresh=0.7,
61            rpn_fg_iou_thresh=0.7,
62            rpn_bg_iou_thresh=0.3,
63            # Box parameters
64            box_score_thresh=0.05,
65            box_nms_thresh=0.5,
66            box_detections_per_img=100,
67            box_fg_iou_thresh=0.5,
68            box_bg_iou_thresh=0.5,
69        )
70
71        return model
72
73    def forward(self, images: torch.Tensor, targets=None):
74        """
75        Forward pass through Mask R-CNN.
76
77        During training: returns losses
78        During inference: returns predictions
79        """
80        return self.model(images, targets)
81
82
83class RoIAlignExplained:
84    """
85    Understand RoIAlign - the key improvement over RoIPool.
86
87    RoIPool: Quantization causes misalignment
88    RoIAlign: Bilinear interpolation preserves spatial accuracy
89    """
90
91    @staticmethod
92    def roi_align_visualization():
93        """
94        RoIAlign samples at regular intervals using bilinear interpolation.
95        """
96        # Simulated feature map
97        feature_map = torch.randn(1, 256, 50, 50)
98
99        # Region of Interest (batch_idx, x1, y1, x2, y2)
100        # Coordinates can be fractional!
101        rois = torch.tensor([
102            [0, 10.3, 15.7, 30.8, 45.2],  # Fractional coordinates
103            [0, 5.5, 10.1, 25.9, 35.6]
104        ])
105
106        # Apply RoIAlign
107        roi_align = torchvision.ops.RoIAlign(
108            output_size=(7, 7),
109            spatial_scale=1.0,
110            sampling_ratio=2,  # Sample 2x2 points per bin
111            aligned=True
112        )
113
114        # Extract aligned features
115        pooled_features = roi_align(feature_map, rois)
116        print(f"Pooled features shape: {pooled_features.shape}")
117        # Output: [2, 256, 7, 7]
118
119        return pooled_features
120
121
122class MaskHead(nn.Module):
123    """
124    The mask prediction head of Mask R-CNN.
125    Takes RoI-aligned features and predicts binary masks.
126    """
127
128    def __init__(self, in_channels: int, num_classes: int):
129        super().__init__()
130
131        # 4 convolutional layers
132        self.conv1 = nn.Conv2d(in_channels, 256, 3, padding=1)
133        self.conv2 = nn.Conv2d(256, 256, 3, padding=1)
134        self.conv3 = nn.Conv2d(256, 256, 3, padding=1)
135        self.conv4 = nn.Conv2d(256, 256, 3, padding=1)
136
137        # Upsampling to increase resolution
138        self.deconv = nn.ConvTranspose2d(256, 256, 2, stride=2)
139
140        # Final mask prediction (one mask per class)
141        self.mask_predictor = nn.Conv2d(256, num_classes, 1)
142
143        self.relu = nn.ReLU(inplace=True)
144
145    def forward(self, x: torch.Tensor) -> torch.Tensor:
146        """
147        Args:
148            x: RoI-aligned features [N, C, 14, 14]
149
150        Returns:
151            masks: [N, num_classes, 28, 28]
152        """
153        x = self.relu(self.conv1(x))
154        x = self.relu(self.conv2(x))
155        x = self.relu(self.conv3(x))
156        x = self.relu(self.conv4(x))
157        x = self.relu(self.deconv(x))
158        masks = self.mask_predictor(x)
159        return masks
160
161# Example usage
162custom_maskrcnn = CustomMaskRCNN(num_classes=91)  # COCO classes
163print(f"Model built with {custom_maskrcnn.backbone_name} backbone")

1import torch
2from torch.utils.data import Dataset, DataLoader
3import torchvision.transforms as T
4import json
5import numpy as np
6from PIL import Image
7from pycocotools import mask as mask_utils
8
9class InstanceSegmentationDataset(Dataset):
10    """
11    Custom dataset for instance segmentation.
12    Supports COCO format annotations.
13    """
14
15    def __init__(self, root_dir: str, annotation_file: str,
16                 transforms=None):
17        self.root_dir = root_dir
18        self.transforms = transforms
19
20        # Load COCO format annotations
21        with open(annotation_file, 'r') as f:
22            self.coco_data = json.load(f)
23
24        self.images = {img['id']: img for img in self.coco_data['images']}
25        self.categories = {cat['id']: cat for cat in self.coco_data['categories']}
26
27        # Group annotations by image
28        self.img_to_anns = {}
29        for ann in self.coco_data['annotations']:
30            img_id = ann['image_id']
31            if img_id not in self.img_to_anns:
32                self.img_to_anns[img_id] = []
33            self.img_to_anns[img_id].append(ann)
34
35        self.image_ids = list(self.img_to_anns.keys())
36
37    def __len__(self):
38        return len(self.image_ids)
39
40    def __getitem__(self, idx):
41        img_id = self.image_ids[idx]
42        img_info = self.images[img_id]
43        annotations = self.img_to_anns[img_id]
44
45        # Load image
46        img_path = f"{self.root_dir}/{img_info['file_name']}"
47        image = Image.open(img_path).convert('RGB')
48
49        # Process annotations
50        boxes = []
51        labels = []
52        masks = []
53        areas = []
54        iscrowd = []
55
56        for ann in annotations:
57            # Bounding box (COCO format: [x, y, width, height])
58            x, y, w, h = ann['bbox']
59            boxes.append([x, y, x + w, y + h])
60
61            # Category label
62            labels.append(ann['category_id'])
63
64            # Segmentation mask
65            if 'segmentation' in ann:
66                if isinstance(ann['segmentation'], dict):
67                    # RLE format
68                    mask = mask_utils.decode(ann['segmentation'])
69                else:
70                    # Polygon format
71                    mask = self._polygon_to_mask(
72                        ann['segmentation'],
73                        img_info['height'],
74                        img_info['width']
75                    )
76                masks.append(mask)
77
78            areas.append(ann['area'])
79            iscrowd.append(ann.get('iscrowd', 0))
80
81        # Convert to tensors
82        target = {
83            'boxes': torch.as_tensor(boxes, dtype=torch.float32),
84            'labels': torch.as_tensor(labels, dtype=torch.int64),
85            'masks': torch.as_tensor(np.array(masks), dtype=torch.uint8),
86            'image_id': torch.tensor([img_id]),
87            'area': torch.as_tensor(areas, dtype=torch.float32),
88            'iscrowd': torch.as_tensor(iscrowd, dtype=torch.int64)
89        }
90
91        if self.transforms:
92            image, target = self.transforms(image, target)
93
94        return image, target
95
96    def _polygon_to_mask(self, polygons, height, width):
97        """Convert polygon segmentation to binary mask."""
98        from pycocotools import mask as mask_utils
99
100        rles = mask_utils.frPyObjects(polygons, height, width)
101        rle = mask_utils.merge(rles)
102        mask = mask_utils.decode(rle)
103        return mask
104
105
106class InstanceSegmentationTrainer:
107    """Complete training pipeline for Mask R-CNN."""
108
109    def __init__(self, model, device='cuda'):
110        self.model = model.to(device)
111        self.device = device
112
113        # Only train non-frozen parameters
114        params = [p for p in model.parameters() if p.requires_grad]
115        self.optimizer = torch.optim.SGD(
116            params,
117            lr=0.005,
118            momentum=0.9,
119            weight_decay=0.0005
120        )
121
122        # Learning rate scheduler
123        self.lr_scheduler = torch.optim.lr_scheduler.StepLR(
124            self.optimizer,
125            step_size=3,
126            gamma=0.1
127        )
128
129    def train_one_epoch(self, dataloader, epoch):
130        """Train for one epoch."""
131        self.model.train()
132
133        total_loss = 0
134        loss_components = {
135            'loss_classifier': 0,
136            'loss_box_reg': 0,
137            'loss_mask': 0,
138            'loss_objectness': 0,
139            'loss_rpn_box_reg': 0
140        }
141
142        for batch_idx, (images, targets) in enumerate(dataloader):
143            images = [img.to(self.device) for img in images]
144            targets = [{k: v.to(self.device) for k, v in t.items()}
145                      for t in targets]
146
147            # Forward pass - returns losses during training
148            loss_dict = self.model(images, targets)
149
150            # Sum all losses
151            losses = sum(loss for loss in loss_dict.values())
152
153            # Backward pass
154            self.optimizer.zero_grad()
155            losses.backward()
156
157            # Gradient clipping
158            torch.nn.utils.clip_grad_norm_(
159                self.model.parameters(),
160                max_norm=1.0
161            )
162
163            self.optimizer.step()
164
165            total_loss += losses.item()
166            for k, v in loss_dict.items():
167                loss_components[k] += v.item()
168
169            if batch_idx % 10 == 0:
170                print(f"Epoch {epoch}, Batch {batch_idx}, Loss: {losses.item():.4f}")
171
172        self.lr_scheduler.step()
173
174        # Average losses
175        num_batches = len(dataloader)
176        avg_loss = total_loss / num_batches
177        avg_components = {k: v / num_batches for k, v in loss_components.items()}
178
179        return avg_loss, avg_components
180
181    @torch.no_grad()
182    def evaluate(self, dataloader):
183        """Evaluate model on validation set."""
184        self.model.eval()
185
186        all_predictions = []
187        all_targets = []
188
189        for images, targets in dataloader:
190            images = [img.to(self.device) for img in images]
191
192            # Get predictions
193            predictions = self.model(images)
194
195            all_predictions.extend(predictions)
196            all_targets.extend(targets)
197
198        return all_predictions, all_targets
199
200# Data augmentation for instance segmentation
201def get_transform(train: bool):
202    transforms = []
203    transforms.append(T.ToTensor())
204
205    if train:
206        transforms.append(T.RandomHorizontalFlip(0.5))
207
208    return T.Compose(transforms)

1import numpy as np
2from typing import List, Dict, Tuple
3from collections import defaultdict
4
5class SegmentationMetrics:
6    """
7    Comprehensive metrics for instance segmentation evaluation.
8    """
9
10    @staticmethod
11    def compute_iou(mask1: np.ndarray, mask2: np.ndarray) -> float:
12        """
13        Compute Intersection over Union between two masks.
14
15        IoU = Area of Intersection / Area of Union
16        """
17        intersection = np.logical_and(mask1, mask2).sum()
18        union = np.logical_or(mask1, mask2).sum()
19
20        if union == 0:
21            return 0.0
22
23        return intersection / union
24
25    @staticmethod
26    def compute_dice(mask1: np.ndarray, mask2: np.ndarray) -> float:
27        """
28        Compute Dice coefficient (F1 for segmentation).
29
30        Dice = 2 * Intersection / (|A| + |B|)
31        """
32        intersection = np.logical_and(mask1, mask2).sum()
33        total = mask1.sum() + mask2.sum()
34
35        if total == 0:
36            return 0.0
37
38        return 2 * intersection / total
39
40    def compute_ap_single_class(self,
41                                predictions: List[Dict],
42                                ground_truths: List[Dict],
43                                iou_threshold: float = 0.5) -> float:
44        """
45        Compute Average Precision for a single class.
46
47        Args:
48            predictions: List of {'mask', 'score', 'image_id'}
49            ground_truths: List of {'mask', 'image_id'}
50            iou_threshold: IoU threshold for matching
51
52        Returns:
53            Average Precision value
54        """
55        # Sort predictions by confidence score
56        predictions = sorted(predictions, key=lambda x: x['score'], reverse=True)
57
58        # Group ground truths by image
59        gt_by_image = defaultdict(list)
60        for gt in ground_truths:
61            gt_by_image[gt['image_id']].append(gt)
62
63        # Track which GTs have been matched
64        gt_matched = {gt['image_id']: [False] * len(gts)
65                     for gt['image_id'], gts in gt_by_image.items()}
66
67        tp = np.zeros(len(predictions))
68        fp = np.zeros(len(predictions))
69
70        for pred_idx, pred in enumerate(predictions):
71            img_id = pred['image_id']
72            pred_mask = pred['mask']
73
74            best_iou = 0
75            best_gt_idx = -1
76
77            # Find best matching ground truth
78            for gt_idx, gt in enumerate(gt_by_image.get(img_id, [])):
79                if gt_matched[img_id][gt_idx]:
80                    continue
81
82                iou = self.compute_iou(pred_mask, gt['mask'])
83                if iou > best_iou:
84                    best_iou = iou
85                    best_gt_idx = gt_idx
86
87            if best_iou >= iou_threshold:
88                tp[pred_idx] = 1
89                gt_matched[img_id][best_gt_idx] = True
90            else:
91                fp[pred_idx] = 1
92
93        # Compute precision-recall curve
94        tp_cumsum = np.cumsum(tp)
95        fp_cumsum = np.cumsum(fp)
96
97        total_gt = len(ground_truths)
98
99        recalls = tp_cumsum / total_gt
100        precisions = tp_cumsum / (tp_cumsum + fp_cumsum)
101
102        # Compute AP using 11-point interpolation
103        ap = 0
104        for t in np.arange(0, 1.1, 0.1):
105            mask = recalls >= t
106            if mask.any():
107                ap += precisions[mask].max() / 11
108
109        return ap
110
111    def compute_map(self,
112                    predictions: Dict[int, List[Dict]],
113                    ground_truths: Dict[int, List[Dict]],
114                    iou_thresholds: List[float] = None) -> Dict:
115        """
116        Compute mean Average Precision across classes and IoU thresholds.
117
118        COCO-style evaluation uses IoU thresholds from 0.5 to 0.95.
119        """
120        if iou_thresholds is None:
121            iou_thresholds = [0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95]
122
123        results = {}
124
125        # Get all classes
126        all_classes = set(predictions.keys()) | set(ground_truths.keys())
127
128        ap_per_class = {}
129        for class_id in all_classes:
130            class_preds = predictions.get(class_id, [])
131            class_gts = ground_truths.get(class_id, [])
132
133            if not class_gts:  # No ground truth for this class
134                continue
135
136            ap_at_thresholds = []
137            for iou_thresh in iou_thresholds:
138                ap = self.compute_ap_single_class(
139                    class_preds, class_gts, iou_thresh
140                )
141                ap_at_thresholds.append(ap)
142
143            ap_per_class[class_id] = {
144                'AP@50': ap_at_thresholds[0] if iou_thresholds[0] == 0.5 else None,
145                'AP@75': ap_at_thresholds[5] if len(iou_thresholds) > 5 else None,
146                'AP': np.mean(ap_at_thresholds)  # COCO-style mAP
147            }
148
149        # Compute mean across classes
150        results['mAP'] = np.mean([v['AP'] for v in ap_per_class.values()])
151        results['mAP@50'] = np.mean([v['AP@50'] for v in ap_per_class.values()
152                                    if v['AP@50'] is not None])
153        results['mAP@75'] = np.mean([v['AP@75'] for v in ap_per_class.values()
154                                    if v['AP@75'] is not None])
155        results['per_class'] = ap_per_class
156
157        return results
158
159    @staticmethod
160    def panoptic_quality(pred_segments: np.ndarray,
161                        gt_segments: np.ndarray) -> Dict[str, float]:
162        """
163        Compute Panoptic Quality (PQ) metric.
164
165        PQ = SQ * RQ
166        - SQ (Segmentation Quality): Average IoU of matched segments
167        - RQ (Recognition Quality): F1 of matched segments
168        """
169        # Find unique segment IDs
170        pred_ids = set(np.unique(pred_segments)) - {0}  # Exclude background
171        gt_ids = set(np.unique(gt_segments)) - {0}
172
173        # Match predictions to ground truth
174        matched_pred = set()
175        matched_gt = set()
176        iou_sum = 0
177
178        for pred_id in pred_ids:
179            pred_mask = pred_segments == pred_id
180            best_iou = 0
181            best_gt_id = None
182
183            for gt_id in gt_ids:
184                if gt_id in matched_gt:
185                    continue
186
187                gt_mask = gt_segments == gt_id
188                intersection = np.logical_and(pred_mask, gt_mask).sum()
189                union = np.logical_or(pred_mask, gt_mask).sum()
190                iou = intersection / union if union > 0 else 0
191
192                if iou > best_iou and iou > 0.5:  # Threshold for matching
193                    best_iou = iou
194                    best_gt_id = gt_id
195
196            if best_gt_id is not None:
197                matched_pred.add(pred_id)
198                matched_gt.add(best_gt_id)
199                iou_sum += best_iou
200
201        tp = len(matched_pred)
202        fp = len(pred_ids) - tp
203        fn = len(gt_ids) - tp
204
205        # Segmentation Quality
206        sq = iou_sum / tp if tp > 0 else 0
207
208        # Recognition Quality (F1)
209        rq = tp / (tp + 0.5 * fp + 0.5 * fn) if (tp + fp + fn) > 0 else 0
210
211        # Panoptic Quality
212        pq = sq * rq
213
214        return {
215            'PQ': pq,
216            'SQ': sq,
217            'RQ': rq,
218            'TP': tp,
219            'FP': fp,
220            'FN': fn
221        }
222
223# Example usage
224metrics = SegmentationMetrics()
225
226# Create example masks
227mask1 = np.zeros((100, 100), dtype=bool)
228mask1[20:60, 30:70] = True
229
230mask2 = np.zeros((100, 100), dtype=bool)
231mask2[25:65, 35:75] = True
232
233iou = metrics.compute_iou(mask1, mask2)
234dice = metrics.compute_dice(mask1, mask2)
235print(f"IoU: {iou:.4f}, Dice: {dice:.4f}")