初始化结果张量，确保与MaskSplit输出格式一致
(self, max_iterations, current_image, current_mask)
| 213 | return image, mask |
| 214 | |
| 215 | def initialize_results(self, max_iterations, current_image, current_mask): |
| 216 | """ |
| 217 | 初始化结果张量，确保与MaskSplit输出格式一致 |
| 218 | """ |
| 219 | # 确保维度正确 |
| 220 | assert len(current_image.shape) == 4, "Current image must be 4D [B,H,W,C]" |
| 221 | assert len(current_mask.shape) == 3, "Current mask must be 3D [B,H,W]" |
| 222 | |
| 223 | # 创建结果张量，确保格式一致 |
| 224 | result_images = torch.zeros( |
| 225 | (max_iterations, current_image.shape[1], current_image.shape[2], current_image.shape[3]), |
| 226 | dtype=current_image.dtype, |
| 227 | device=current_image.device |
| 228 | ) # 明确指定 [B,H,W,C] |
| 229 | |
| 230 | result_masks = torch.zeros( |
| 231 | (max_iterations, current_mask.shape[1], current_mask.shape[2]), |
| 232 | dtype=current_mask.dtype, |
| 233 | device=current_mask.device |
| 234 | ) # 明确指定 [B,H,W] |
| 235 | |
| 236 | return result_images, result_masks |
| 237 | |
| 238 | def while_loop_close(self, flow_control, current_image, current_mask, max_iterations, |
| 239 | pass_back=False, iteration_count=0, result_images=None, result_masks=None, |