标准化输入格式 image: 确保是4D tensor [B,H,W,C] mask: 确保是3D tensor [B,H,W]
(self, image, mask)
| 195 | self.collect_contained(child_id, upstream, contained) |
| 196 | |
| 197 | def standardize_input(self, image, mask): |
| 198 | """ |
| 199 | 标准化输入格式 |
| 200 | image: 确保是4D tensor [B,H,W,C] |
| 201 | mask: 确保是3D tensor [B,H,W] |
| 202 | """ |
| 203 | # 处理image |
| 204 | if len(image.shape) == 3: # [H,W,C] -> [1,H,W,C] |
| 205 | image = image.unsqueeze(0) |
| 206 | assert len(image.shape) == 4, f"Image must be 4D [B,H,W,C], got shape {image.shape}" |
| 207 | |
| 208 | # 处理mask |
| 209 | if len(mask.shape) == 2: # [H,W] -> [1,H,W] |
| 210 | mask = mask.unsqueeze(0) |
| 211 | assert len(mask.shape) == 3, f"Mask must be 3D [B,H,W], got shape {mask.shape}" |
| 212 | |
| 213 | return image, mask |
| 214 | |
| 215 | def initialize_results(self, max_iterations, current_image, current_mask): |
| 216 | """ |