Coverage for bioimageio/core/proc_ops.py: 77%

1import collections.abc 

2from abc import ABC, abstractmethod 

3from dataclasses import InitVar, dataclass, field 

4from typing import ( 

5 Collection, 

6 Literal, 

7 Mapping, 

8 Optional, 

9 Sequence, 

10 Set, 

11 Tuple, 

12 Union, 

13) 

14 

15import numpy as np 

16import xarray as xr 

17from typing_extensions import Self, assert_never 

18 

19from bioimageio.spec.model import v0_4, v0_5 

20 

21from ._op_base import BlockedOperator, Operator 

22from .axis import AxisId, PerAxis 

23from .block import Block 

24from .common import DTypeStr, MemberId 

25from .sample import Sample, SampleBlock, SampleBlockWithOrigin 

26from .stat_calculators import StatsCalculator 

27from .stat_measures import ( 

28 DatasetMean, 

29 DatasetMeasure, 

30 DatasetPercentile, 

31 DatasetStd, 

32 MeanMeasure, 

33 Measure, 

34 MeasureValue, 

35 SampleMean, 

36 SampleQuantile, 

37 SampleStd, 

38 Stat, 

39 StdMeasure, 

40) 

41from .tensor import Tensor 

42 

43 

44def _convert_axis_ids( 

45 axes: v0_4.AxesInCZYX, 

46 mode: Literal["per_sample", "per_dataset"], 

47) -> Tuple[AxisId, ...]: 

48 if not isinstance(axes, str): 

49 return tuple(axes) 

50 

51 if mode == "per_sample": 

52 ret = [] 

53 elif mode == "per_dataset": 

54 ret = [AxisId("b")] 

55 else: 

56 assert_never(mode) 

57 

58 ret.extend([AxisId(a) for a in axes]) 

59 return tuple(ret) 

60 

61 

62@dataclass 

63class _SimpleOperator(BlockedOperator, ABC): 

64 input: MemberId 

65 output: MemberId 

66 

67 @property 

68 def required_measures(self) -> Collection[Measure]: 

69 return set() 

70 

71 @abstractmethod 

72 def get_output_shape(self, input_shape: PerAxis[int]) -> PerAxis[int]: ... 

73 

74 def __call__(self, sample: Union[Sample, SampleBlock]) -> None: 

75 if self.input not in sample.members: 

76 return 

77 

78 input_tensor = sample.members[self.input] 

79 output_tensor = self._apply(input_tensor, sample.stat) 

80 

81 if self.output in sample.members: 

82 assert ( 

83 sample.members[self.output].tagged_shape == output_tensor.tagged_shape 

84 ) 

85 

86 if isinstance(sample, Sample): 

87 sample.members[self.output] = output_tensor 

88 elif isinstance(sample, SampleBlock): 

89 b = sample.blocks[self.input] 

90 sample.blocks[self.output] = Block( 

91 sample_shape=self.get_output_shape(sample.shape[self.input]), 

92 data=output_tensor, 

93 inner_slice=b.inner_slice, 

94 halo=b.halo, 

95 block_index=b.block_index, 

96 blocks_in_sample=b.blocks_in_sample, 

97 ) 

98 else: 

99 assert_never(sample) 

100 

101 @abstractmethod 

102 def _apply(self, input: Tensor, stat: Stat) -> Tensor: ... 

103 

104 

105@dataclass 

106class AddKnownDatasetStats(BlockedOperator): 

107 dataset_stats: Mapping[DatasetMeasure, MeasureValue] 

108 

109 @property 

110 def required_measures(self) -> Set[Measure]: 

111 return set() 

112 

113 def __call__(self, sample: Union[Sample, SampleBlock]) -> None: 

114 sample.stat.update(self.dataset_stats.items()) 

115 

116 

117# @dataclass 

118# class UpdateStats(Operator): 

119# """Calculates sample and/or dataset measures""" 

120 

121# measures: Union[Sequence[Measure], Set[Measure], Mapping[Measure, MeasureValue]] 

122# """sample and dataset `measuers` to be calculated by this operator. Initial/fixed 

123# dataset measure values may be given, see `keep_updating_dataset_stats` for details. 

124# """ 

125# keep_updating_dataset_stats: Optional[bool] = None 

126# """indicates if operator calls should keep updating dataset statistics or not 

127 

128# default (None): if `measures` is a `Mapping` (i.e. initial measure values are 

129# given) no further updates to dataset statistics is conducted, otherwise (w.o. 

130# initial measure values) dataset statistics are updated by each processed sample. 

131# """ 

132# _keep_updating_dataset_stats: bool = field(init=False) 

133# _stats_calculator: StatsCalculator = field(init=False) 

134 

135# @property 

136# def required_measures(self) -> Set[Measure]: 

137# return set() 

138 

139# def __post_init__(self): 

140# self._stats_calculator = StatsCalculator(self.measures) 

141# if self.keep_updating_dataset_stats is None: 

142# self._keep_updating_dataset_stats = not isinstance(self.measures, collections.abc.Mapping) 

143# else: 

144# self._keep_updating_dataset_stats = self.keep_updating_dataset_stats 

145 

146# def __call__(self, sample_block: SampleBlockWithOrigin> None: 

147# if self._keep_updating_dataset_stats: 

148# sample.stat.update(self._stats_calculator.update_and_get_all(sample)) 

149# else: 

150# sample.stat.update(self._stats_calculator.skip_update_and_get_all(sample)) 

151 

152 

153@dataclass 

154class UpdateStats(Operator): 

155 """Calculates sample and/or dataset measures""" 

156 

157 stats_calculator: StatsCalculator 

158 """`StatsCalculator` to be used by this operator.""" 

159 keep_updating_initial_dataset_stats: bool = False 

160 """indicates if operator calls should keep updating initial dataset statistics or not; 

161 if the `stats_calculator` was not provided with any initial dataset statistics, 

162 these are always updated with every new sample. 

163 """ 

164 _keep_updating_dataset_stats: bool = field(init=False) 

165 

166 @property 

167 def required_measures(self) -> Set[Measure]: 

168 return set() 

169 

170 def __post_init__(self): 

171 self._keep_updating_dataset_stats = ( 

172 self.keep_updating_initial_dataset_stats 

173 or not self.stats_calculator.has_dataset_measures 

174 ) 

175 

176 def __call__(self, sample: Union[Sample, SampleBlockWithOrigin]) -> None: 

177 if isinstance(sample, SampleBlockWithOrigin): 

178 # update stats with whole sample on first block 

179 if sample.block_index != 0: 

180 return 

181 

182 origin = sample.origin 

183 else: 

184 origin = sample 

185 

186 if self._keep_updating_dataset_stats: 

187 sample.stat.update(self.stats_calculator.update_and_get_all(origin)) 

188 else: 

189 sample.stat.update(self.stats_calculator.skip_update_and_get_all(origin)) 

190 

191 

192@dataclass 

193class Binarize(_SimpleOperator): 

194 """'output = tensor > threshold'.""" 

195 

196 threshold: Union[float, Sequence[float]] 

197 axis: Optional[AxisId] = None 

198 

199 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

200 return input > self.threshold 

201 

202 def get_output_shape( 

203 self, input_shape: Mapping[AxisId, int] 

204 ) -> Mapping[AxisId, int]: 

205 return input_shape 

206 

207 @classmethod 

208 def from_proc_descr( 

209 cls, descr: Union[v0_4.BinarizeDescr, v0_5.BinarizeDescr], member_id: MemberId 

210 ) -> Self: 

211 if isinstance(descr.kwargs, (v0_4.BinarizeKwargs, v0_5.BinarizeKwargs)): 

212 return cls( 

213 input=member_id, output=member_id, threshold=descr.kwargs.threshold 

214 ) 

215 elif isinstance(descr.kwargs, v0_5.BinarizeAlongAxisKwargs): 

216 return cls( 

217 input=member_id, 

218 output=member_id, 

219 threshold=descr.kwargs.threshold, 

220 axis=descr.kwargs.axis, 

221 ) 

222 else: 

223 assert_never(descr.kwargs) 

224 

225 

226@dataclass 

227class Clip(_SimpleOperator): 

228 min: Optional[float] = None 

229 """minimum value for clipping""" 

230 max: Optional[float] = None 

231 """maximum value for clipping""" 

232 

233 def __post_init__(self): 

234 assert self.min is not None or self.max is not None, "missing min or max value" 

235 assert ( 

236 self.min is None or self.max is None or self.min < self.max 

237 ), f"expected min < max, but {self.min} !< {self.max}" 

238 

239 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

240 return input.clip(self.min, self.max) 

241 

242 def get_output_shape( 

243 self, input_shape: Mapping[AxisId, int] 

244 ) -> Mapping[AxisId, int]: 

245 return input_shape 

246 

247 @classmethod 

248 def from_proc_descr( 

249 cls, descr: Union[v0_4.ClipDescr, v0_5.ClipDescr], member_id: MemberId 

250 ) -> Self: 

251 return cls( 

252 input=member_id, 

253 output=member_id, 

254 min=descr.kwargs.min, 

255 max=descr.kwargs.max, 

256 ) 

257 

258 

259@dataclass 

260class EnsureDtype(_SimpleOperator): 

261 dtype: DTypeStr 

262 

263 @classmethod 

264 def from_proc_descr(cls, descr: v0_5.EnsureDtypeDescr, member_id: MemberId): 

265 return cls(input=member_id, output=member_id, dtype=descr.kwargs.dtype) 

266 

267 def get_descr(self): 

268 return v0_5.EnsureDtypeDescr(kwargs=v0_5.EnsureDtypeKwargs(dtype=self.dtype)) 

269 

270 def get_output_shape( 

271 self, input_shape: Mapping[AxisId, int] 

272 ) -> Mapping[AxisId, int]: 

273 return input_shape 

274 

275 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

276 return input.astype(self.dtype) 

277 

278 

279@dataclass 

280class ScaleLinear(_SimpleOperator): 

281 gain: Union[float, xr.DataArray] = 1.0 

282 """multiplicative factor""" 

283 

284 offset: Union[float, xr.DataArray] = 0.0 

285 """additive term""" 

286 

287 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

288 return input * self.gain + self.offset 

289 

290 def get_output_shape( 

291 self, input_shape: Mapping[AxisId, int] 

292 ) -> Mapping[AxisId, int]: 

293 return input_shape 

294 

295 @classmethod 

296 def from_proc_descr( 

297 cls, 

298 descr: Union[v0_4.ScaleLinearDescr, v0_5.ScaleLinearDescr], 

299 member_id: MemberId, 

300 ) -> Self: 

301 kwargs = descr.kwargs 

302 if isinstance(kwargs, v0_5.ScaleLinearAlongAxisKwargs): 

303 axis = kwargs.axis 

304 elif isinstance(kwargs, (v0_4.ScaleLinearKwargs, v0_5.ScaleLinearKwargs)): 

305 axis = None 

306 else: 

307 assert_never(kwargs) 

308 

309 if axis: 

310 gain = xr.DataArray(np.atleast_1d(kwargs.gain), dims=axis) 

311 offset = xr.DataArray(np.atleast_1d(kwargs.offset), dims=axis) 

312 else: 

313 assert ( 

314 isinstance(kwargs.gain, (float, int)) or len(kwargs.gain) == 1 

315 ), kwargs.gain 

316 gain = ( 

317 kwargs.gain if isinstance(kwargs.gain, (float, int)) else kwargs.gain[0] 

318 ) 

319 assert isinstance(kwargs.offset, (float, int)) or len(kwargs.offset) == 1 

320 offset = ( 

321 kwargs.offset 

322 if isinstance(kwargs.offset, (float, int)) 

323 else kwargs.offset[0] 

324 ) 

325 

326 return cls(input=member_id, output=member_id, gain=gain, offset=offset) 

327 

328 

329@dataclass 

330class ScaleMeanVariance(_SimpleOperator): 

331 axes: Optional[Sequence[AxisId]] = None 

332 reference_tensor: Optional[MemberId] = None 

333 eps: float = 1e-6 

334 mean: Union[SampleMean, DatasetMean] = field(init=False) 

335 std: Union[SampleStd, DatasetStd] = field(init=False) 

336 ref_mean: Union[SampleMean, DatasetMean] = field(init=False) 

337 ref_std: Union[SampleStd, DatasetStd] = field(init=False) 

338 

339 @property 

340 def required_measures(self): 

341 return {self.mean, self.std, self.ref_mean, self.ref_std} 

342 

343 def __post_init__(self): 

344 axes = None if self.axes is None else tuple(self.axes) 

345 ref_tensor = self.reference_tensor or self.input 

346 if axes is None or AxisId("batch") not in axes: 

347 Mean = SampleMean 

348 Std = SampleStd 

349 else: 

350 Mean = DatasetMean 

351 Std = DatasetStd 

352 

353 self.mean = Mean(member_id=self.input, axes=axes) 

354 self.std = Std(member_id=self.input, axes=axes) 

355 self.ref_mean = Mean(member_id=ref_tensor, axes=axes) 

356 self.ref_std = Std(member_id=ref_tensor, axes=axes) 

357 

358 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

359 mean = stat[self.mean] 

360 std = stat[self.std] + self.eps 

361 ref_mean = stat[self.ref_mean] 

362 ref_std = stat[self.ref_std] + self.eps 

363 return (input - mean) / std * ref_std + ref_mean 

364 

365 def get_output_shape( 

366 self, input_shape: Mapping[AxisId, int] 

367 ) -> Mapping[AxisId, int]: 

368 return input_shape 

369 

370 @classmethod 

371 def from_proc_descr( 

372 cls, 

373 descr: Union[v0_4.ScaleMeanVarianceDescr, v0_5.ScaleMeanVarianceDescr], 

374 member_id: MemberId, 

375 ) -> Self: 

376 kwargs = descr.kwargs 

377 _, axes = _get_axes(descr.kwargs) 

378 

379 return cls( 

380 input=member_id, 

381 output=member_id, 

382 reference_tensor=MemberId(str(kwargs.reference_tensor)), 

383 axes=axes, 

384 eps=kwargs.eps, 

385 ) 

386 

387 

388def _get_axes( 

389 kwargs: Union[ 

390 v0_4.ZeroMeanUnitVarianceKwargs, 

391 v0_5.ZeroMeanUnitVarianceKwargs, 

392 v0_4.ScaleRangeKwargs, 

393 v0_5.ScaleRangeKwargs, 

394 v0_4.ScaleMeanVarianceKwargs, 

395 v0_5.ScaleMeanVarianceKwargs, 

396 ], 

397) -> Tuple[bool, Optional[Tuple[AxisId, ...]]]: 

398 if kwargs.axes is None: 

399 return True, None 

400 elif isinstance(kwargs.axes, str): 

401 axes = _convert_axis_ids(kwargs.axes, kwargs["mode"]) 

402 return AxisId("b") in axes, axes 

403 elif isinstance(kwargs.axes, collections.abc.Sequence): 

404 axes = tuple(kwargs.axes) 

405 return AxisId("batch") in axes, axes 

406 else: 

407 assert_never(kwargs.axes) 

408 

409 

410@dataclass 

411class ScaleRange(_SimpleOperator): 

412 lower_percentile: InitVar[Optional[Union[SampleQuantile, DatasetPercentile]]] = None 

413 upper_percentile: InitVar[Optional[Union[SampleQuantile, DatasetPercentile]]] = None 

414 lower: Union[SampleQuantile, DatasetPercentile] = field(init=False) 

415 upper: Union[SampleQuantile, DatasetPercentile] = field(init=False) 

416 

417 eps: float = 1e-6 

418 

419 def __post_init__( 

420 self, 

421 lower_percentile: Optional[Union[SampleQuantile, DatasetPercentile]], 

422 upper_percentile: Optional[Union[SampleQuantile, DatasetPercentile]], 

423 ): 

424 if lower_percentile is None: 

425 tid = self.input if upper_percentile is None else upper_percentile.member_id 

426 self.lower = DatasetPercentile(q=0.0, member_id=tid) 

427 else: 

428 self.lower = lower_percentile 

429 

430 if upper_percentile is None: 

431 self.upper = DatasetPercentile(q=1.0, member_id=self.lower.member_id) 

432 else: 

433 self.upper = upper_percentile 

434 

435 assert self.lower.member_id == self.upper.member_id 

436 assert self.lower.q < self.upper.q 

437 assert self.lower.axes == self.upper.axes 

438 

439 @property 

440 def required_measures(self): 

441 return {self.lower, self.upper} 

442 

443 def get_output_shape( 

444 self, input_shape: Mapping[AxisId, int] 

445 ) -> Mapping[AxisId, int]: 

446 return input_shape 

447 

448 @classmethod 

449 def from_proc_descr( 

450 cls, 

451 descr: Union[v0_4.ScaleRangeDescr, v0_5.ScaleRangeDescr], 

452 member_id: MemberId, 

453 ): 

454 kwargs = descr.kwargs 

455 ref_tensor = ( 

456 member_id 

457 if kwargs.reference_tensor is None 

458 else MemberId(str(kwargs.reference_tensor)) 

459 ) 

460 dataset_mode, axes = _get_axes(descr.kwargs) 

461 if dataset_mode: 

462 Percentile = DatasetPercentile 

463 else: 

464 Percentile = SampleQuantile 

465 

466 return cls( 

467 input=member_id, 

468 output=member_id, 

469 lower_percentile=Percentile( 

470 q=kwargs.min_percentile / 100, axes=axes, member_id=ref_tensor 

471 ), 

472 upper_percentile=Percentile( 

473 q=kwargs.max_percentile / 100, axes=axes, member_id=ref_tensor 

474 ), 

475 ) 

476 

477 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

478 lower = stat[self.lower] 

479 upper = stat[self.upper] 

480 return (input - lower) / (upper - lower + self.eps) 

481 

482 def get_descr(self): 

483 assert self.lower.axes == self.upper.axes 

484 assert self.lower.member_id == self.upper.member_id 

485 

486 return v0_5.ScaleRangeDescr( 

487 kwargs=v0_5.ScaleRangeKwargs( 

488 axes=self.lower.axes, 

489 min_percentile=self.lower.q * 100, 

490 max_percentile=self.upper.q * 100, 

491 eps=self.eps, 

492 reference_tensor=self.lower.member_id, 

493 ) 

494 ) 

495 

496 

497@dataclass 

498class Sigmoid(_SimpleOperator): 

499 """1 / (1 + e^(-input)).""" 

500 

501 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

502 return Tensor(1.0 / (1.0 + np.exp(-input)), dims=input.dims) 

503 

504 @property 

505 def required_measures(self) -> Collection[Measure]: 

506 return {} 

507 

508 def get_output_shape( 

509 self, input_shape: Mapping[AxisId, int] 

510 ) -> Mapping[AxisId, int]: 

511 return input_shape 

512 

513 @classmethod 

514 def from_proc_descr( 

515 cls, descr: Union[v0_4.SigmoidDescr, v0_5.SigmoidDescr], member_id: MemberId 

516 ) -> Self: 

517 assert isinstance(descr, (v0_4.SigmoidDescr, v0_5.SigmoidDescr)) 

518 return cls(input=member_id, output=member_id) 

519 

520 def get_descr(self): 

521 return v0_5.SigmoidDescr() 

522 

523 

524@dataclass 

525class ZeroMeanUnitVariance(_SimpleOperator): 

526 """normalize to zero mean, unit variance.""" 

527 

528 mean: MeanMeasure 

529 std: StdMeasure 

530 

531 eps: float = 1e-6 

532 

533 def __post_init__(self): 

534 assert self.mean.axes == self.std.axes 

535 

536 @property 

537 def required_measures(self) -> Set[Union[MeanMeasure, StdMeasure]]: 

538 return {self.mean, self.std} 

539 

540 def get_output_shape( 

541 self, input_shape: Mapping[AxisId, int] 

542 ) -> Mapping[AxisId, int]: 

543 return input_shape 

544 

545 @classmethod 

546 def from_proc_descr( 

547 cls, 

548 descr: Union[v0_4.ZeroMeanUnitVarianceDescr, v0_5.ZeroMeanUnitVarianceDescr], 

549 member_id: MemberId, 

550 ): 

551 dataset_mode, axes = _get_axes(descr.kwargs) 

552 

553 if dataset_mode: 

554 Mean = DatasetMean 

555 Std = DatasetStd 

556 else: 

557 Mean = SampleMean 

558 Std = SampleStd 

559 

560 return cls( 

561 input=member_id, 

562 output=member_id, 

563 mean=Mean(axes=axes, member_id=member_id), 

564 std=Std(axes=axes, member_id=member_id), 

565 ) 

566 

567 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

568 mean = stat[self.mean] 

569 std = stat[self.std] 

570 return (input - mean) / (std + self.eps) 

571 

572 def get_descr(self): 

573 return v0_5.ZeroMeanUnitVarianceDescr( 

574 kwargs=v0_5.ZeroMeanUnitVarianceKwargs(axes=self.mean.axes, eps=self.eps) 

575 ) 

576 

577 

578@dataclass 

579class FixedZeroMeanUnitVariance(_SimpleOperator): 

580 """normalize to zero mean, unit variance with precomputed values.""" 

581 

582 mean: Union[float, xr.DataArray] 

583 std: Union[float, xr.DataArray] 

584 

585 eps: float = 1e-6 

586 

587 def __post_init__(self): 

588 assert ( 

589 isinstance(self.mean, (int, float)) 

590 or isinstance(self.std, (int, float)) 

591 or self.mean.dims == self.std.dims 

592 ) 

593 

594 def get_output_shape( 

595 self, input_shape: Mapping[AxisId, int] 

596 ) -> Mapping[AxisId, int]: 

597 return input_shape 

598 

599 @classmethod 

600 def from_proc_descr( 

601 cls, 

602 descr: v0_5.FixedZeroMeanUnitVarianceDescr, 

603 member_id: MemberId, 

604 ) -> Self: 

605 if isinstance(descr.kwargs, v0_5.FixedZeroMeanUnitVarianceKwargs): 

606 dims = None 

607 elif isinstance(descr.kwargs, v0_5.FixedZeroMeanUnitVarianceAlongAxisKwargs): 

608 dims = (descr.kwargs.axis,) 

609 else: 

610 assert_never(descr.kwargs) 

611 

612 return cls( 

613 input=member_id, 

614 output=member_id, 

615 mean=xr.DataArray(descr.kwargs.mean, dims=dims), 

616 std=xr.DataArray(descr.kwargs.std, dims=dims), 

617 ) 

618 

619 def get_descr(self): 

620 if isinstance(self.mean, (int, float)): 

621 assert isinstance(self.std, (int, float)) 

622 kwargs = v0_5.FixedZeroMeanUnitVarianceKwargs(mean=self.mean, std=self.std) 

623 else: 

624 assert isinstance(self.std, xr.DataArray) 

625 assert len(self.mean.dims) == 1 

626 kwargs = v0_5.FixedZeroMeanUnitVarianceAlongAxisKwargs( 

627 axis=AxisId(str(self.mean.dims[0])), 

628 mean=list(self.mean), 

629 std=list(self.std), 

630 ) 

631 

632 return v0_5.FixedZeroMeanUnitVarianceDescr(kwargs=kwargs) 

633 

634 def _apply(self, input: Tensor, stat: Stat) -> Tensor: 

635 return (input - self.mean) / (self.std + self.eps) 

636 

637 

638ProcDescr = Union[ 

639 v0_4.PreprocessingDescr, 

640 v0_4.PostprocessingDescr, 

641 v0_5.PreprocessingDescr, 

642 v0_5.PostprocessingDescr, 

643] 

644 

645Processing = Union[ 

646 AddKnownDatasetStats, 

647 Binarize, 

648 Clip, 

649 EnsureDtype, 

650 FixedZeroMeanUnitVariance, 

651 ScaleLinear, 

652 ScaleMeanVariance, 

653 ScaleRange, 

654 Sigmoid, 

655 UpdateStats, 

656 ZeroMeanUnitVariance, 

657] 

658 

659 

660def get_proc_class(proc_spec: ProcDescr): 

661 if isinstance(proc_spec, (v0_4.BinarizeDescr, v0_5.BinarizeDescr)): 

662 return Binarize 

663 elif isinstance(proc_spec, (v0_4.ClipDescr, v0_5.ClipDescr)): 

664 return Clip 

665 elif isinstance(proc_spec, v0_5.EnsureDtypeDescr): 

666 return EnsureDtype 

667 elif isinstance(proc_spec, v0_5.FixedZeroMeanUnitVarianceDescr): 

668 return FixedZeroMeanUnitVariance 

669 elif isinstance(proc_spec, (v0_4.ScaleLinearDescr, v0_5.ScaleLinearDescr)): 

670 return ScaleLinear 

671 elif isinstance( 

672 proc_spec, (v0_4.ScaleMeanVarianceDescr, v0_5.ScaleMeanVarianceDescr) 

673 ): 

674 return ScaleMeanVariance 

675 elif isinstance(proc_spec, (v0_4.ScaleRangeDescr, v0_5.ScaleRangeDescr)): 

676 return ScaleRange 

677 elif isinstance(proc_spec, (v0_4.SigmoidDescr, v0_5.SigmoidDescr)): 

678 return Sigmoid 

679 elif ( 

680 isinstance(proc_spec, v0_4.ZeroMeanUnitVarianceDescr) 

681 and proc_spec.kwargs.mode == "fixed" 

682 ): 

683 return FixedZeroMeanUnitVariance 

684 elif isinstance( 

685 proc_spec, 

686 (v0_4.ZeroMeanUnitVarianceDescr, v0_5.ZeroMeanUnitVarianceDescr), 

687 ): 

688 return ZeroMeanUnitVariance 

689 else: 

690 assert_never(proc_spec)