| import torch | |
| def custom_collate(batch): | |
| images = [item['image'] for item in batch] | |
| cls_labels = [item['cls_label'] for item in batch] | |
| # Some items might not have masks | |
| masks = [item['mask'] for item in batch] | |
| patient_ids = [item['patient_id'] for item in batch] | |
| # Stack only non-None masks; keep None as placeholder | |
| # or pad with dummy masks for consistent shape | |
| collated_batch = { | |
| 'image': torch.stack(images), | |
| 'patient_id': patient_ids | |
| } | |
| # --------- Handle cls_label (optional) ---------- | |
| if any(label is not None for label in cls_labels): | |
| # Use first non-None label as reference for shape | |
| for label in cls_labels: | |
| if label is not None: | |
| ref_shape = label.shape | |
| break | |
| dummy_label = torch.zeros(ref_shape, dtype=torch.long) | |
| cls_labels = [lbl if lbl is not None else dummy_label for lbl in cls_labels] | |
| collated_batch['cls_label'] = torch.stack(cls_labels) | |
| collated_batch['has_cls_label'] = torch.tensor([lbl is not None for lbl in cls_labels]) | |
| else: | |
| collated_batch['cls_label'] = None | |
| collated_batch['has_cls_label'] = None | |
| # Check if at least one sample has a mask | |
| if any(mask is not None for mask in masks): | |
| # Replace None with a dummy zero mask matching the shape | |
| # ref_shape = masks[0].shape if masks[0] is not None else masks[1].shape | |
| for mask in masks: | |
| if mask is not None: | |
| ref_shape = mask.shape | |
| break | |
| dummy = torch.zeros(ref_shape, dtype=torch.long) | |
| masks = [m if m is not None else dummy for m in masks] | |
| collated_batch['mask'] = torch.stack(masks) | |
| collated_batch['has_mask'] = torch.tensor([m is not None for m in masks]) | |
| else: | |
| collated_batch['mask'] = None | |
| collated_batch['has_mask'] = None | |
| return collated_batch | |