Coverage for ezstitcher/core/pipeline_factories.py: 100%

1""" 

2Pipeline factory system for the EZStitcher pipeline architecture. 

3 

4This module contains the AutoPipelineFactory class that creates pre-configured pipelines 

5for all common workflows, leveraging specialized steps to reduce boilerplate code. 

6 

7The AutoPipelineFactory uses a unified approach that handles 2D multichannel, z-stack per plane stitch, 

8and z-stack projection stitch with a single implementation, simplifying the pipeline architecture. 

9""" 

10 

11from typing import List, Optional, Union, Dict, Any 

12from pathlib import Path 

13 

14from .pipeline import Pipeline 

15from .steps import Step, PositionGenerationStep, ImageStitchingStep, ZFlatStep, FocusStep, CompositeStep, NormStep 

16from .image_processor import ImageProcessor as IP 

17 

18 

19class AutoPipelineFactory: 

20 """ 

21 Unified factory for creating pipelines for all common use cases. 

22 

23 This factory handles all types of stitching workflows with a single implementation: 

24 - 2D multichannel stitching 

25 - Z-stack per plane stitching 

26 - Z-stack projection stitching 

27 

28 It automatically configures the appropriate steps based on the input parameters, 

29 with no need to differentiate between different types of pipelines. 

30 """ 

31 

32 def __init__( 

33 self, 

34 input_dir: Union[str, Path], 

35 output_dir: Optional[Union[str, Path]] = None, 

36 normalize: bool = True, 

37 normalization_params: Optional[Dict[str, Any]] = None, 

38 well_filter: Optional[List[str]] = None, 

39 flatten_z: bool = False, 

40 z_method: str = "max", 

41 channel_weights: Optional[Union[List[float], Dict[str, float]]] = None, 

42 ): 

43 """ 

44 Initialize with pipeline parameters. 

45 

46 Args: 

47 input_dir: Input directory containing images 

48 output_dir: Output directory for stitched images 

49 normalize: Whether to include normalization 

50 normalization_params: Parameters for normalization 

51 well_filter: Wells to process 

52 flatten_z: Whether to flatten Z-stacks (if Z-stacks are present) 

53 z_method: Z-stack processing method: 

54 - "max", "mean", "median" for projection methods 

55 - "combined", "laplacian", "tenengrad" for focus detection methods 

56 channel_weights: Weights for channel compositing (for reference image only). 

57 Should be a list with length equal to the number of channels. 

58 """ 

59 self.input_dir = Path(input_dir) 

60 # Only use output_dir if explicitly provided, otherwise let the pipeline handle it 

61 self.output_dir = Path(output_dir) if output_dir else None 

62 

63 self.normalize = normalize 

64 # Default normalization parameters 

65 self.normalization_params = normalization_params or { 

66 'low_percentile': 1.0, 

67 'high_percentile': 99.0 

68 } 

69 self.well_filter = well_filter 

70 self.flatten_z = flatten_z 

71 self.z_method = z_method 

72 self.channel_weights = channel_weights 

73 

74 # Determine if z_method is a focus method or projection method 

75 self.focus_methods = ["combined", "laplacian", "tenengrad", "normalized_variance", "fft"] 

76 self.is_focus_method = self.z_method in self.focus_methods 

77 

78 def create_pipelines(self) -> List[Pipeline]: 

79 """ 

80 Create pipeline configuration based on parameters. 

81 

82 This method creates two pipelines: 

83 1. Position generation pipeline - Creates position files for stitching 

84 2. Image assembly pipeline - Stitches images using the position files 

85 """ 

86 # Create position generation pipeline 

87 pos_pipeline = Pipeline( 

88 input_dir=self.input_dir, 

89 steps=[ 

90 # Always include Z-flattening for position generation 

91 ZFlatStep(method="max"), 

92 

93 # Include normalization if enabled 

94 NormStep(**self.normalization_params) if self.normalize else None, 

95 

96 # Always include channel compositing for reference image 

97 CompositeStep(weights=self.channel_weights), 

98 

99 # Always include position generation 

100 PositionGenerationStep() 

101 ], 

102 well_filter=self.well_filter, 

103 name="Position Generation Pipeline" 

104 ) 

105 

106 # Get the position generation step's output directory 

107 positions_dir = pos_pipeline.steps[-1].output_dir 

108 

109 # Create assembly pipeline 

110 assembly_pipeline = Pipeline( 

111 input_dir=self.input_dir, 

112 output_dir=self.output_dir, 

113 steps=[ 

114 # Include normalization if enabled (create new instance) 

115 NormStep(**self.normalization_params) if self.normalize else None, 

116 

117 # Include Z-flattening for assembly if enabled 

118 (FocusStep( 

119 focus_options={'metric': self.z_method} 

120 ) if self.is_focus_method else 

121 ZFlatStep( 

122 method=self.z_method 

123 )) if self.flatten_z else None, 

124 

125 # Always include image stitching with explicit positions directory 

126 ImageStitchingStep(positions_dir=positions_dir) 

127 ], 

128 well_filter=self.well_filter, 

129 name="Image Assembly Pipeline" 

130 ) 

131 

132 # Return both pipelines 

133 return [pos_pipeline, assembly_pipeline]