Coverage for src/bioimageio/core/_prediction_pipeline.py: 75%

1import warnings 

2from abc import ABC, abstractmethod 

3from itertools import chain 

4from types import MappingProxyType 

5from typing import ( 

6 Any, 

7 Iterable, 

8 List, 

9 Literal, 

10 Mapping, 

11 NamedTuple, 

12 Optional, 

13 Sequence, 

14 Set, 

15 Tuple, 

16 TypeVar, 

17 Union, 

18) 

19 

20from loguru import logger 

21from tqdm import tqdm 

22from typing_extensions import assert_never 

23 

24from bioimageio.spec import load_model_description 

25from bioimageio.spec.model import AnyModelDescr, v0_4, v0_5 

26 

27from ._model_adapter import ModelAdapter 

28from ._op_base import BlockwiseOperator, SamplewiseOperator 

29from .axis import AxisId, PerAxis 

30from .backends import create_model_adapter 

31from .common import ( 

32 BlocksizeParameter, 

33 Halo, 

34 MemberId, 

35 PerMember, 

36 SampleId, 

37 SupportedWeightsFormat, 

38) 

39from .digest_spec import ( 

40 get_block_transform, 

41 get_input_halo, 

42 get_member_ids, 

43) 

44from .proc_ops import Processing 

45from .proc_setup import setup_pre_and_postprocessing 

46from .sample import Sample, SampleBlock 

47from .stat_measures import Measure, MeasureValue, Stat 

48from .tensor import Tensor 

49 

50Predict_IO = TypeVar( 

51 "Predict_IO", 

52 Sample, 

53 Iterable[Sample], 

54) 

55 

56 

57class IntermediatePrediction(NamedTuple): 

58 """Represents an intermediate prediction of a sample with blocking, including the predicted sample so far and the last predicted block. 

59 

60 The final `IntermediatePrediction` in a sequence holds the complete predicted (and postprocessed if applicable) sample.""" 

61 

62 sample: Sample 

63 last_block: SampleBlock 

64 

65 

66class _PredictionPipelineBase(ABC): 

67 def __init__( 

68 self, 

69 model_descr: AnyModelDescr, 

70 *, 

71 default_blocksize_parameter: BlocksizeParameter, 

72 default_batch_size: int, 

73 preceding_prediction_pipelines: Optional[ 

74 Sequence[Union["PredictionPipeline", "RemotePredictionPipeline"]] 

75 ], 

76 ) -> None: 

77 super().__init__() 

78 self._model_descr = model_descr 

79 self._default_blocksize_parameter = default_blocksize_parameter 

80 self._default_batch_size = default_batch_size 

81 # TODO: Improve parallelization of blockwise predictions with preceding prediction pipelines. 

82 self._preceding_prediction_pipelines = preceding_prediction_pipelines 

83 if isinstance(model_descr, v0_4.ModelDescr): 

84 self._default_output_halo: PerMember[PerAxis[Halo]] = {} 

85 self._default_input_halo: PerMember[PerAxis[Halo]] = {} 

86 self._block_transform = None 

87 else: 

88 self._default_output_halo = { 

89 t.id: { 

90 a.id: Halo(a.halo, a.halo) 

91 for a in t.axes 

92 if isinstance(a, v0_5.WithHalo) 

93 } 

94 for t in model_descr.outputs 

95 } 

96 self._default_input_halo = get_input_halo( 

97 model_descr, self._default_output_halo 

98 ) 

99 self._block_transform = get_block_transform(model_descr) 

100 

101 self.pad_mode = ( 

102 {} 

103 if isinstance(model_descr, v0_4.ModelDescr) 

104 else { 

105 descr.id: descr.pad or v0_5.SymmetricPadding() 

106 for descr in model_descr.inputs 

107 } 

108 ) 

109 

110 self._input_ids = tuple(get_member_ids(model_descr.inputs)) 

111 self._output_ids = tuple(get_member_ids(model_descr.outputs)) 

112 

113 @property 

114 def input_ids(self) -> Sequence[MemberId]: 

115 return self._input_ids 

116 

117 @property 

118 def output_ids(self) -> Sequence[MemberId]: 

119 return self._output_ids 

120 

121 @property 

122 def model_descr(self) -> AnyModelDescr: 

123 return self._model_descr 

124 

125 @property 

126 def model_description(self) -> AnyModelDescr: 

127 return self._model_descr 

128 

129 def _get_preceding_prediction_pipelines_for_sample( 

130 self, sample: Sample 

131 ) -> Sequence["PredictionPipeline | RemotePredictionPipeline"]: 

132 """Get preceding prediction pipelines for a sample based on the sample's input member ids.""" 

133 if not self._preceding_prediction_pipelines: 

134 return () 

135 

136 required_inputs = set(self.input_ids) 

137 sample_members = set(sample.members.keys()) 

138 preceding_pipelines: List["PredictionPipeline | RemotePredictionPipeline"] = [] 

139 for pp in self._preceding_prediction_pipelines[::-1]: 

140 preceding_pipelines.insert(0, pp) 

141 sample_members.update(pp.output_ids) 

142 required_inputs.update(pp.input_ids) 

143 required_inputs.difference_update(sample_members) 

144 

145 if not required_inputs: 

146 return preceding_pipelines 

147 

148 raise KeyError( 

149 f"Sample is missing required inputs {required_inputs} for the prediction pipeline or its preceding pipelines." 

150 ) 

151 

152 def predict_sample_without_blocking( 

153 self, 

154 sample: Sample, 

155 skip_preprocessing: bool = False, 

156 skip_postprocessing: bool = False, 

157 skip_input_padding: bool = False, 

158 skip_output_cropping: bool = False, 

159 ) -> Sample: 

160 """Predict a whole sample at once. 

161 

162 Note: 

163 The sample's tensor shapes have to match the model's input tensor description. 

164 If that is not the case, consider `predict_sample_with_blocking` 

165 

166 Args: 

167 sample: input sample 

168 skip_preprocessing: if `True`, skip all preprocessing steps (except for any preceding prediction pipeline). 

169 skip_postprocessing: if `True`, skip all postprocessing steps (except for any preceding prediction pipeline). 

170 skip_input_padding: if `True`, skip padding the input sample according to the model's (optional) output halos. 

171 skip_output_cropping: if `True`, skip cropping any output halos from the model output. 

172 """ 

173 for pp in self._get_preceding_prediction_pipelines_for_sample(sample): 

174 sample = pp._predict_sample_without_blocking_impl( 

175 sample, 

176 skip_input_padding=skip_input_padding, 

177 skip_output_cropping=skip_output_cropping, 

178 ) 

179 

180 return self._predict_sample_without_blocking_impl( 

181 sample, 

182 skip_preprocessing=skip_preprocessing, 

183 skip_postprocessing=skip_postprocessing, 

184 skip_input_padding=skip_input_padding, 

185 skip_output_cropping=skip_output_cropping, 

186 ) 

187 

188 @abstractmethod 

189 def _predict_sample_without_blocking_impl( 

190 self, 

191 sample: Sample, 

192 skip_preprocessing: bool = False, 

193 skip_postprocessing: bool = False, 

194 skip_input_padding: bool = False, 

195 skip_output_cropping: bool = False, 

196 ) -> Sample: 

197 """Predict a whole sample at once. 

198 

199 Note: 

200 The sample's tensor shapes have to match the model's input tensor description. 

201 If that is not the case, consider `predict_sample_with_blocking` 

202 

203 Args: 

204 sample: input sample 

205 skip_preprocessing: if `True`, skip all preprocessing steps. 

206 skip_postprocessing: if `True`, skip all postprocessing steps. 

207 skip_input_padding: if `True`, skip padding the input sample according to the model's (optional) output halos. 

208 skip_output_cropping: if `True`, skip cropping any output halos from the model output. 

209 """ 

210 

211 def predict_sample_with_blocking( 

212 self, 

213 sample: Sample, 

214 skip_preprocessing: bool = False, 

215 skip_postprocessing: bool = False, 

216 ns: Optional[ 

217 Union[ 

218 v0_5.ParameterizedSize_N, 

219 Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize_N], 

220 ] 

221 ] = None, 

222 batch_size: Optional[int] = None, 

223 ) -> Sample: 

224 """Predict a sample by predicting sample blocks. 

225 

226 Note: For fixed/known blocksizes use `predict_sample_with_fixed_blocking`. 

227 

228 Args: 

229 sample: The sample to predict on. 

230 skip_preprocessing: If `True`, skip all preprocessing steps. 

231 skip_postprocessing: If `True`, skip all postprocessing steps. 

232 ns: Block size parameter(s) allows scaling the model's default input block size. 

233 Blocksize parameters are only applied to parameterized input axes, all other axis sizes are fixed/derived or (for output axes) data dependent. 

234 Unapplicable blocksize parameters are ignored. 

235 batch_size: Batch size to use for prediction. 

236 """ 

237 

238 output = None 

239 for output in self.predict_sample_with_blocking_yield_intermediates( 

240 sample, 

241 skip_preprocessing=skip_preprocessing, 

242 skip_postprocessing=skip_postprocessing, 

243 ns=ns, 

244 batch_size=batch_size, 

245 )[1]: 

246 pass 

247 

248 assert output is not None, ( 

249 "No blocks were predicted, cannot return final sample." 

250 ) 

251 return output.sample 

252 

253 def predict_sample_with_fixed_blocking( 

254 self, 

255 sample: Sample, 

256 input_block_shape: PerMember[PerAxis[int]], 

257 skip_preprocessing: bool = False, 

258 skip_postprocessing: bool = False, 

259 ) -> Sample: 

260 """Predict `sample` with given `input_block_shape`. 

261 

262 Note: 

263 - `input_block_shape` is expected to be a valid input shape for the model. 

264 - Use `predict_sample_with_blocking` if you want to control block sizes via generic block size parameters rather than fixed block shapes. 

265 

266 Args: 

267 sample: The sample to predict on. 

268 input_block_shape: Mapping of input member id to mapping of axis id to block size for that axis. 

269 skip_preprocessing: If `True`, skip all preprocessing steps. 

270 skip_postprocessing: If `True`, skip all postprocessing steps. 

271 """ 

272 intermediate = None 

273 for ( 

274 intermediate 

275 ) in self._predict_sample_with_fixed_blocking_yield_intermediates_impl( 

276 sample, 

277 input_block_shape=input_block_shape, 

278 skip_preprocessing=skip_preprocessing, 

279 skip_postprocessing=skip_postprocessing, 

280 )[1]: 

281 pass 

282 

283 assert intermediate is not None, ( 

284 "No blocks were predicted, cannot return final sample." 

285 ) 

286 return intermediate.sample 

287 

288 def predict_sample_with_blocking_yield_intermediates( 

289 self, 

290 sample: Sample, 

291 skip_preprocessing: bool = False, 

292 skip_postprocessing: bool = False, 

293 ns: Optional[ 

294 Union[ 

295 v0_5.ParameterizedSize_N, 

296 Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize_N], 

297 ] 

298 ] = None, 

299 batch_size: Optional[int] = None, 

300 ) -> Tuple[int, Iterable[IntermediatePrediction]]: 

301 """Predict `sample` by predicting sample blocks and yield intermediate predictions if no samplewise postprocessing is included. 

302 Also yields intermediate predictions if there are preceding prediction pipelines (model inputs depend on another model's outputs). 

303 For preceding prediction pipelines `ns` and `batch_size` are shared, but pre- and postprocessing are never skipped in preceding pipelines. 

304 

305 Returns: 

306 Tuple of number of prediction steps and an iterator of predicted intermediate samples with the last predicted block, 

307 All samples, but the last one, are intermediate samples with more and more blocks predicted. 

308 In case samplewise postprocessing needs to be applied, no intermediate results are yielded, but only the final sample after all blocks are predicted and postprocessed. 

309 In case of preceding prediction pipelines (model inputs depend on another model's outputs), intermediate results initially do not include the final output tensors at all. 

310 """ 

311 

312 total_prediction_steps = 0 

313 iterable_intermediates = () 

314 for pp in self._get_preceding_prediction_pipelines_for_sample(sample): 

315 pp_steps, pp_intermediates = ( 

316 pp.predict_sample_with_blocking_yield_intermediates( 

317 sample, 

318 ns=ns, 

319 batch_size=batch_size, 

320 skip_preprocessing=False, 

321 skip_postprocessing=False, 

322 ) 

323 ) 

324 total_prediction_steps += pp_steps 

325 iterable_intermediates = chain(iterable_intermediates, pp_intermediates) 

326 

327 if isinstance(self._model_descr, v0_4.ModelDescr): 

328 raise NotImplementedError( 

329 "`predict_sample_with_blocking` not implemented for v0_4.ModelDescr" 

330 + f" {self._model_descr.name}." 

331 + " Consider using `predict_sample_with_fixed_blocking`" 

332 ) 

333 

334 ns = ns or self._default_blocksize_parameter 

335 if isinstance(ns, int): 

336 ns = { 

337 (ipt.id, a.id): ns 

338 for ipt in self._model_descr.inputs 

339 for a in ipt.axes 

340 if isinstance(a.size, v0_5.ParameterizedSize) 

341 } 

342 input_block_shape = self._model_descr.get_tensor_sizes( 

343 ns, batch_size or self._default_batch_size 

344 ).inputs 

345 

346 steps, intermediates = ( 

347 self._predict_sample_with_fixed_blocking_yield_intermediates_impl( 

348 sample, 

349 input_block_shape=input_block_shape, 

350 skip_preprocessing=skip_preprocessing, 

351 skip_postprocessing=skip_postprocessing, 

352 ) 

353 ) 

354 total_prediction_steps += steps 

355 iterable_intermediates = chain(iterable_intermediates, intermediates) 

356 return total_prediction_steps, iterable_intermediates 

357 

358 def predict_sample_with_fixed_blocking_yield_intermediates( 

359 self, 

360 sample: Sample, 

361 input_block_shape: PerMember[PerAxis[int]], 

362 *, 

363 skip_preprocessing: bool = False, 

364 skip_postprocessing: bool = False, 

365 fill_value: float = float("nan"), 

366 ) -> Tuple[int, Iterable[IntermediatePrediction]]: 

367 """Predict `sample` by predicting sample blocks of `input_block_shape` and yield intermediate predictions if no samplewise postprocessing is included. 

368 Also yields intermediate predictions if there are preceding prediction pipelines (model inputs depend on another model's outputs). 

369 For preceding prediction pipelines `input_block_shape` and `fill_value` are shared, but pre- and postprocessing are never skipped in preceding pipelines. 

370 

371 Returns: 

372 Tuple of number of prediction steps and an iterator of predicted intermediate samples with the last predicted block, 

373 All samples, but the last one, are intermediate samples with more and more blocks predicted. 

374 In case samplewise postprocessing needs to be applied, no intermediate results are yielded, but only the final sample after all blocks are predicted and postprocessed. 

375 In case of preceding prediction pipelines (model inputs depend on another model's outputs), intermediate results initially do not include the final output tensors at all. 

376 """ 

377 total_prediction_steps = 0 

378 iterable_intermediates = () 

379 for pp in self._get_preceding_prediction_pipelines_for_sample(sample): 

380 pp_steps, pp_intermediates = ( 

381 pp._predict_sample_with_fixed_blocking_yield_intermediates_impl( 

382 sample, 

383 input_block_shape=input_block_shape, 

384 skip_preprocessing=False, 

385 skip_postprocessing=False, 

386 fill_value=fill_value, 

387 ) 

388 ) 

389 total_prediction_steps += pp_steps 

390 iterable_intermediates = chain(iterable_intermediates, pp_intermediates) 

391 

392 pp_steps, pp_intermediates = ( 

393 self._predict_sample_with_fixed_blocking_yield_intermediates_impl( 

394 sample, 

395 input_block_shape=input_block_shape, 

396 skip_preprocessing=skip_preprocessing, 

397 skip_postprocessing=skip_postprocessing, 

398 fill_value=fill_value, 

399 ) 

400 ) 

401 total_prediction_steps += pp_steps 

402 iterable_intermediates = chain(iterable_intermediates, pp_intermediates) 

403 return total_prediction_steps, iterable_intermediates 

404 

405 @abstractmethod 

406 def _predict_sample_with_fixed_blocking_yield_intermediates_impl( 

407 self, 

408 sample: Sample, 

409 input_block_shape: PerMember[PerAxis[int]], 

410 *, 

411 skip_preprocessing: bool = False, 

412 skip_postprocessing: bool = False, 

413 fill_value: float = float("nan"), 

414 ) -> Tuple[int, Iterable[IntermediatePrediction]]: ... 

415 

416 @abstractmethod 

417 def predict_sample_block( 

418 self, 

419 sample_block: SampleBlock, 

420 skip_preprocessing: bool = False, 

421 skip_postprocessing: bool = False, 

422 ) -> SampleBlock: 

423 """Predict a single sample block. 

424 

425 Note that this does not apply samplewise preprocessing or postprocessing steps, but only blockwise ones. 

426 

427 Args: 

428 sample_block: The sample block to predict on. 

429 skip_preprocessing: If `True`, skip blockwise preprocessing steps. 

430 skip_postprocessing: If `True`, skip blockwise postprocessing steps. 

431 """ 

432 

433 

434class PredictionPipeline(_PredictionPipelineBase): 

435 """ 

436 Represents model computation including preprocessing and postprocessing 

437 Note: Ideally use the `PredictionPipeline` in a with statement 

438 (as a context manager). 

439 """ 

440 

441 def __init__( 

442 self, 

443 *, 

444 name: str, 

445 model_description: AnyModelDescr, 

446 preprocessing: List[Processing], 

447 postprocessing: List[Processing], 

448 model_adapter: ModelAdapter, 

449 default_blocksize_parameter: BlocksizeParameter = 10, 

450 default_batch_size: int = 1, 

451 preceding_prediction_pipelines: Optional[ 

452 Sequence[Union["PredictionPipeline", "RemotePredictionPipeline"]] 

453 ] = None, 

454 ) -> None: 

455 """Consider using `create_prediction_pipeline` to create a `PredictionPipeline` with sensible defaults.""" 

456 super().__init__( 

457 model_descr=model_description, 

458 default_blocksize_parameter=default_blocksize_parameter, 

459 default_batch_size=default_batch_size, 

460 preceding_prediction_pipelines=preceding_prediction_pipelines, 

461 ) 

462 

463 if model_description.run_mode: 

464 warnings.warn( 

465 f"Not yet implemented inference for run mode '{model_description.run_mode.name}'" 

466 ) 

467 

468 self.name = name 

469 # split preprocessing into samplewise and blockwise. samplewise preprocessing is all preprocessing up to including the last samplewise operator, blockwise preprocessing are the remaining blockwise operators. 

470 # I.e. some samplewise preprocessing may be a blockwise op (at some point followed by a samplewise op). 

471 self._samplewise_preprocessing: List[ 

472 Union[SamplewiseOperator, BlockwiseOperator] 

473 ] = [] 

474 self._blockwise_preprocessing: List[BlockwiseOperator] = [] 

475 for op in preprocessing[::-1]: 

476 if isinstance(op, BlockwiseOperator) and not self._samplewise_preprocessing: 

477 self._blockwise_preprocessing.insert(0, op) 

478 else: 

479 self._samplewise_preprocessing.insert(0, op) 

480 # split postprocessing analougly, but here we start blockwise and switch to samplewise at the first samplewise operator. 

481 self._blockwise_postprocessing: List[BlockwiseOperator] = [] 

482 self._samplewise_postprocessing: List[ 

483 Union[BlockwiseOperator, SamplewiseOperator] 

484 ] = [] 

485 for op in postprocessing: 

486 if ( 

487 isinstance(op, BlockwiseOperator) 

488 and not self._samplewise_postprocessing 

489 ): 

490 self._blockwise_postprocessing.append(op) 

491 else: 

492 self._samplewise_postprocessing.append(op) 

493 

494 self._adapter = model_adapter 

495 

496 def __enter__(self): 

497 self.load() 

498 return self 

499 

500 def __exit__(self, exc_type, exc_val, exc_tb): # type: ignore 

501 self.unload() 

502 return False 

503 

504 @property 

505 def has_non_blockwise_preprocessing(self) -> bool: 

506 """`True` if any preprocessing operators in the pipeline are not applicable blockwise.""" 

507 return bool(self._samplewise_preprocessing) 

508 

509 @property 

510 def has_non_blockwise_postprocessing(self) -> bool: 

511 """`True` if any postprocessing operators in the pipeline are not applicable blockwise.""" 

512 return bool(self._samplewise_postprocessing) 

513 

514 def _raise_for_non_blockwise_processing( 

515 self, proc_type: Literal["preprocessing", "postprocessing"] 

516 ): 

517 ops = ( 

518 self._samplewise_preprocessing 

519 if proc_type == "preprocessing" 

520 else self._samplewise_postprocessing 

521 ) 

522 non_blockwise = [ 

523 op.__class__.__name__ for op in ops if not isinstance(op, BlockwiseOperator) 

524 ] 

525 if non_blockwise: 

526 raise NotImplementedError( 

527 f"Blockwise {proc_type} for {non_blockwise} not implemented." 

528 ) 

529 

530 def raise_for_non_blockwise_preprocessing(self): 

531 """ 

532 Raises: 

533 NotImplementedError: if there are any non-blockwise preprocessing operators in the pipeline 

534 """ 

535 self._raise_for_non_blockwise_processing("preprocessing") 

536 

537 def raise_for_non_blockwise_postprocessing(self): 

538 """ 

539 Raises: 

540 NotImplementedError: if there are any non-blockwise postprocessing operators in the pipeline 

541 """ 

542 self._raise_for_non_blockwise_processing("postprocessing") 

543 

544 def predict_sample_block( 

545 self, 

546 sample_block: SampleBlock, 

547 skip_preprocessing: bool = False, 

548 skip_postprocessing: bool = False, 

549 ) -> SampleBlock: 

550 if isinstance(self._model_descr, v0_4.ModelDescr): 

551 raise NotImplementedError( 

552 f"predict_sample_block not implemented for model {self._model_descr.format_version}" 

553 ) 

554 else: 

555 assert self._block_transform is not None 

556 

557 if not skip_preprocessing: 

558 self._apply_blockwise_preprocessing(sample_block) 

559 

560 output_meta = sample_block.get_transformed_meta(self._block_transform) 

561 local_output = self._adapter.forward(sample_block.members) 

562 

563 output = output_meta.with_data( 

564 {k: v for k, v in local_output.items() if v is not None}, 

565 stat=sample_block.stat, 

566 ) 

567 if not skip_postprocessing: 

568 self._apply_blockwise_postprocessing(output) 

569 

570 return output 

571 

572 def _predict_sample_without_blocking_impl( 

573 self, 

574 sample: Sample, 

575 skip_preprocessing: bool = False, 

576 skip_postprocessing: bool = False, 

577 skip_input_padding: bool = False, 

578 skip_output_cropping: bool = False, 

579 ) -> Sample: 

580 if not skip_input_padding: 

581 sample = sample.pad(pad_width=self._default_input_halo, mode=self.pad_mode) 

582 

583 if not skip_preprocessing: 

584 self.apply_preprocessing(sample) 

585 

586 output = Sample( 

587 members={ 

588 k: v 

589 for k, v in self._adapter.forward(sample.members).items() 

590 if v is not None 

591 }, 

592 stat=sample.stat, 

593 id=sample.id, 

594 ) 

595 if not skip_postprocessing: 

596 self.apply_postprocessing(output) 

597 

598 if not skip_output_cropping: 

599 output.members = { 

600 m: t 

601 if m not in self._default_output_halo 

602 else t[ 

603 { 

604 a: slice(h.left, None if h.right == 0 else -h.right) 

605 for a, h in self._default_output_halo[m].items() 

606 } 

607 ] 

608 for m, t in output.members.items() 

609 } 

610 

611 return output 

612 

613 def get_output_sample_id(self, input_sample_id: SampleId): 

614 warnings.warn( 

615 "`PredictionPipeline.get_output_sample_id()` is deprecated and will be" 

616 + " removed soon. Output sample id is equal to input sample id, hence this" 

617 + " function is not needed." 

618 ) 

619 return input_sample_id 

620 

621 def predict_sample_with_blocking( 

622 self, 

623 sample: Sample, 

624 skip_preprocessing: bool = False, 

625 skip_postprocessing: bool = False, 

626 ns: Optional[ 

627 Union[ 

628 v0_5.ParameterizedSize_N, 

629 Mapping[Tuple[MemberId, AxisId], v0_5.ParameterizedSize_N], 

630 ] 

631 ] = None, 

632 batch_size: Optional[int] = None, 

633 ) -> Sample: 

634 output = None 

635 for output in self.predict_sample_with_blocking_yield_intermediates( 

636 sample, 

637 skip_preprocessing=skip_preprocessing, 

638 skip_postprocessing=skip_postprocessing, 

639 ns=ns, 

640 batch_size=batch_size, 

641 )[1]: 

642 pass 

643 

644 assert output is not None, ( 

645 "No blocks were predicted, cannot return final sample." 

646 ) 

647 return output.sample 

648 

649 def _predict_sample_with_fixed_blocking_yield_intermediates_impl( 

650 self, 

651 sample: Sample, 

652 input_block_shape: Mapping[MemberId, Mapping[AxisId, int]], 

653 *, 

654 skip_preprocessing: bool = False, 

655 skip_postprocessing: bool = False, 

656 fill_value: float = float("nan"), 

657 ) -> Tuple[int, Iterable[IntermediatePrediction]]: 

658 """Predict `sample` with given `input_block_shape` and yield the full sample with intermediate results. 

659 

660 Note: 

661 - `input_block_shape` is expected to be a valid input shape for the model. 

662 - Use `predict_sample_with_blocking` if you want to control block sizes via generic block size parameters 

663 rather than fixed block shapes. 

664 - Postprocessing may only be complete for the final sample (if samplewise postprocessing steps are included 

665 in the pipeline), intermediate samples may have some (blockwise applicable) postprocessing steps applied. 

666 

667 Args: 

668 sample: The sample to predict on. 

669 input_block_shape: Mapping of input member id to mapping of axis id to block size for that axis. 

670 skip_preprocessing: If `True`, skip all preprocessing steps. 

671 skip_postprocessing: If `True`, skip all postprocessing steps. 

672 

673 Returns: 

674 Tuple of number of blocks and an iterable of predicted intermediate samples with the last predicted block, 

675 All samples, but the last one, are intermediate samples with more and more blocks predicted. 

676 """ 

677 if not skip_preprocessing: 

678 self._apply_samplewise_preprocessing(sample) 

679 

680 n_blocks, input_blocks = sample.split_into_blocks( 

681 input_block_shape, 

682 halo=self._default_input_halo, 

683 pad_mode=self.pad_mode, 

684 ) 

685 logger.info( 

686 "split sample shape {} into {} blocks of {}.", 

687 {k: dict(v) for k, v in sample.shape.items()}, 

688 n_blocks, 

689 {k: dict(v) for k, v in input_block_shape.items()}, 

690 ) 

691 

692 def _predict_blocks(): 

693 predicted_sample = None 

694 for i, b in enumerate( 

695 tqdm( 

696 input_blocks, 

697 desc=f"predict sample {sample.id or ''} with {self._model_descr.id or self._model_descr.name}", 

698 unit="block", 

699 unit_divisor=1, 

700 total=n_blocks, 

701 ) 

702 ): 

703 if not skip_preprocessing: 

704 self._apply_blockwise_preprocessing(b) 

705 

706 predicted_block = self.predict_sample_block( 

707 b, skip_preprocessing=True, skip_postprocessing=True 

708 ) 

709 

710 if not skip_postprocessing: 

711 self._apply_blockwise_postprocessing(predicted_block) 

712 

713 if predicted_sample is None: 

714 predicted_sample = Sample.from_blocks( 

715 [predicted_block], fill_value=fill_value 

716 ) 

717 else: 

718 predicted_sample.set_block(predicted_block) 

719 

720 if not skip_postprocessing and i == n_blocks - 1: 

721 self._apply_samplewise_postprocessing(predicted_sample) 

722 

723 yield IntermediatePrediction(predicted_sample, predicted_block) 

724 

725 return n_blocks, _predict_blocks() 

726 

727 def _apply_samplewise_preprocessing(self, sample: Sample, /) -> None: 

728 """Apply preprocessing operators up to and including the last samplewise operator in-place. 

729 

730 Note: This skips all blockwise preprocessing steps after the last samplewise operator. 

731 """ 

732 if isinstance(sample, SampleBlock): 

733 self.raise_for_non_blockwise_preprocessing() 

734 

735 for op in self._samplewise_preprocessing: 

736 op(sample) 

737 

738 def _apply_blockwise_preprocessing( 

739 self, sample_block: Union[Sample, SampleBlock], / 

740 ) -> None: 

741 """Apply blockwise preprocessing operators in-place. 

742 

743 Note: This skips all preprocessing operators up to and including the last samplewise one. 

744 """ 

745 for op in self._blockwise_preprocessing: 

746 op(sample_block) 

747 

748 def apply_preprocessing(self, sample: Union[Sample, SampleBlock]) -> None: 

749 """Apply preprocessing in-place, also may updates sample stats""" 

750 

751 if isinstance(sample, Sample): 

752 self._apply_samplewise_preprocessing(sample) 

753 else: 

754 self.raise_for_non_blockwise_preprocessing() 

755 

756 self._apply_blockwise_preprocessing(sample) 

757 

758 def _apply_blockwise_postprocessing( 

759 self, sample_block: Union[Sample, SampleBlock], / 

760 ) -> None: 

761 """Apply in-place blockwise postprocessing operators 

762 

763 Note: This does not apply all postprocessing operators from the first samplewise one onwards. 

764 """ 

765 for op in self._blockwise_postprocessing: 

766 op(sample_block) 

767 

768 def _apply_samplewise_postprocessing(self, sample: Sample, /) -> None: 

769 """Apply in-place postprocessing operators starting from and including the first samplewise operator. 

770 

771 Note: This skips all blockwise postprocessing steps before the first samplewise one. 

772 """ 

773 if isinstance(sample, SampleBlock): 

774 self.raise_for_non_blockwise_postprocessing() 

775 

776 for op in self._samplewise_postprocessing: 

777 op(sample) 

778 

779 def apply_postprocessing(self, sample: Union[Sample, SampleBlock]) -> None: 

780 """apply postprocessing in-place, also may updates samples stats""" 

781 self._apply_blockwise_postprocessing(sample) 

782 if isinstance(sample, Sample): 

783 self._apply_samplewise_postprocessing(sample) 

784 else: 

785 self.raise_for_non_blockwise_postprocessing() 

786 

787 def load(self): 

788 """Prepare prediction pipeline for use. 

789 

790 Reusable model adapters may be loaded and unloaded multiple times, but currently not all model adapters 

791 cleanly unload and reload. 

792 

793 Note: 

794 For some model adapters loading is currently part of the constructor making them unusable after unloading. 

795 """ 

796 self._adapter.load() 

797 

798 def unload(self): 

799 """Free any device memory in use. 

800 

801 Note: 

802 Currently prediction pipeline becomes unusable after unloading.""" 

803 self._adapter.unload() 

804 

805 def close(self): 

806 """Permanently close the prediction pipeline and free any device memory in use. 

807 This makes the prediction pipeline unusable afterwards.""" 

808 self.unload() 

809 

810 

811class RemotePredictionPipeline(_PredictionPipelineBase): 

812 """Abstract base class for fully remote prediction pipelines. 

813 

814 Note: A ("local") `PredictionPipeline` may also use a `RemoteModelAdapter` for remote model inference, but it may 

815 still apply local preprocessing and postprocessing steps. 

816 In contrast, a `RemotePredictionPipeline` is designed for the case where all steps including preprocessing and 

817 postprocessing are performed remotely. 

818 """ 

819 

820 def __init__( 

821 self, 

822 model_descr: AnyModelDescr, 

823 *, 

824 server: str, 

825 default_blocksize_parameter: BlocksizeParameter, 

826 default_batch_size: int, 

827 preceding_prediction_pipelines: Optional[ 

828 Sequence[Union["PredictionPipeline", "RemotePredictionPipeline"]] 

829 ] = None, 

830 ) -> None: 

831 super().__init__( 

832 model_descr, 

833 default_blocksize_parameter=default_blocksize_parameter, 

834 default_batch_size=default_batch_size, 

835 preceding_prediction_pipelines=preceding_prediction_pipelines, 

836 ) 

837 self._server = server 

838 

839 @property 

840 def server(self) -> str: 

841 return self._server 

842 

843 

844def create_prediction_pipeline( 

845 bioimageio_model: AnyModelDescr, 

846 *, 

847 devices: Optional[Sequence[str]] = None, 

848 weight_format: Optional[SupportedWeightsFormat] = None, 

849 weights_format: Optional[SupportedWeightsFormat] = None, 

850 dataset_for_initial_statistics: Iterable[Union[Sample, Sequence[Tensor]]] = tuple(), 

851 keep_updating_initial_dataset_statistics: bool = False, 

852 fixed_dataset_statistics: Mapping[Measure, MeasureValue] = MappingProxyType({}), 

853 model_adapter: Optional[ModelAdapter] = None, 

854 ns: Optional[BlocksizeParameter] = None, 

855 default_blocksize_parameter: BlocksizeParameter = 10, # TODO: default to None and find smart blocksize params per axis to reduce overlap of blocks with large halo 

856 preceding_prediction_pipelines: Optional[ 

857 Sequence[Union[PredictionPipeline, RemotePredictionPipeline]] 

858 ] = None, 

859 **deprecated_kwargs: Any, 

860) -> PredictionPipeline: 

861 """ 

862 Creates prediction pipeline which includes: 

863 * computation of input statistics 

864 * preprocessing 

865 * model prediction 

866 * computation of output statistics 

867 * postprocessing 

868 

869 Args: 

870 bioimageio_model: A bioimageio model description. 

871 devices: (optional) 

872 weight_format: deprecated in favor of **weights_format** 

873 weights_format: (optional) Use a specific **weights_format** rather than 

874 choosing one automatically. 

875 A corresponding `bioimageio.core.model_adapters.ModelAdapter` will be 

876 created to run inference with the **bioimageio_model**. 

877 dataset_for_initial_statistics: (optional) If preprocessing steps require input 

878 dataset statistics, **dataset_for_initial_statistics** allows you to 

879 specifcy a dataset from which these statistics are computed. 

880 keep_updating_initial_dataset_statistics: (optional) Set to `True` if you want 

881 to update dataset statistics with each processed sample. 

882 fixed_dataset_statistics: (optional) Precomputed dataset (and optionally sample) statistics. 

883 Any included sample statistics will not be calculated on the fly and it is the callers 

884 responsibility to use samples with the corresponding statistics availble in `sample.stat`. 

885 model_adapter: (optional) Allows you to use a custom **model_adapter** instead 

886 of creating one according to the present/selected **weights_format**. 

887 ns: deprecated in favor of **default_blocksize_parameter** 

888 default_blocksize_parameter: Allows to control the default block size for 

889 blockwise predictions, see `BlocksizeParameter`. 

890 preceding_prediction_pipelines: (optional) If the model has inputs that are 

891 outputs of other models (input field 'output_of'), you can provide a sequence 

892 of preceding prediction pipelines. The prediction pipeline will then automatically 

893 use the outputs of those preceding pipelines as inputs for the current model. 

894 If no preceding prediction pipelines for a model are provided, prediction pipelines using the 

895 same devices and weight format as for the current model will be created for any required preceding models. 

896 """ 

897 weights_format = weight_format or weights_format 

898 del weight_format 

899 default_blocksize_parameter = ns or default_blocksize_parameter 

900 del ns 

901 if deprecated_kwargs: 

902 warnings.warn( 

903 f"deprecated create_prediction_pipeline kwargs: {set(deprecated_kwargs)}" 

904 ) 

905 

906 model_adapter = model_adapter or create_model_adapter( 

907 model_description=bioimageio_model, 

908 devices=devices, 

909 weight_format_priority_order=weights_format and (weights_format,), 

910 ) 

911 

912 input_ids = get_member_ids(bioimageio_model.inputs) 

913 

914 def dataset(): 

915 common_stat: Stat = {} 

916 for i, x in enumerate(dataset_for_initial_statistics): 

917 if isinstance(x, Sample): 

918 yield x 

919 else: 

920 yield Sample(members=dict(zip(input_ids, x)), stat=common_stat, id=i)