标准化输入格式 images: 确保是4D tensor [B,H,W,C] masks: 确保是3D tensor [B,H,W] 如果images是单张图片，会扩展到与masks相同的批次大小
(self, images, masks)
| 30 | CATEGORY = "CyberEveLoop🐰" |
| 31 | |
| 32 | def standardize_input(self, images, masks): |
| 33 | """ |
| 34 | 标准化输入格式 |
| 35 | images: 确保是4D tensor [B,H,W,C] |
| 36 | masks: 确保是3D tensor [B,H,W] |
| 37 | 如果images是单张图片，会扩展到与masks相同的批次大小 |
| 38 | """ |
| 39 | # 处理masks（先处理masks以获取批次大小） |
| 40 | if isinstance(masks, list): |
| 41 | masks = torch.cat(masks, dim=0) |
| 42 | if len(masks.shape) == 2: # [H,W] -> [1,H,W] |
| 43 | masks = masks.unsqueeze(0) |
| 44 | assert len(masks.shape) == 3, f"Masks must be 3D [B,H,W], got shape {masks.shape}" |
| 45 | |
| 46 | # 处理images |
| 47 | if isinstance(images, list): |
| 48 | images = torch.cat(images, dim=0) |
| 49 | if len(images.shape) == 3: # [H,W,C] -> [1,H,W,C] |
| 50 | images = images.unsqueeze(0) |
| 51 | assert len(images.shape) == 4, f"Images must be 4D [B,H,W,C], got shape {images.shape}" |
| 52 | |
| 53 | # 检查是否需要扩展images |
| 54 | if images.shape[0] == 1 and masks.shape[0] > 1: |
| 55 | print(f"Expanding single image to match mask batch size: {masks.shape[0]}") |
| 56 | images = images.expand(masks.shape[0], -1, -1, -1) |
| 57 | |
| 58 | # 确保batch维度相同 |
| 59 | assert images.shape[0] == masks.shape[0], \ |
| 60 | f"Batch size mismatch: images {images.shape[0]} vs masks {masks.shape[0]}" |
| 61 | |
| 62 | return images, masks |
| 63 | |
| 64 | |
| 65 | def resize_to_match(self, image, target_shape): |