Skills › Scientific Agent Skills › openpiv › runner.py
runner.py
268 lines · skills/openpiv/scripts/runner.py at 5f2bc03
1"""CLI for OpenPIV processing of a single image pair.
3Verified against openpiv 0.26.1. Writes vectors.txt, params.npz, and vector_field.png
4into the output directory.
5"""
7import argparse
8from pathlib import Path
9from importlib.metadata import version
11import matplotlib
13# Select a non-interactive backend before pyplot is imported anywhere.
14matplotlib.use("Agg")
16import matplotlib.pyplot as plt # noqa: E402
17import numpy as np # noqa: E402
18from openpiv import filters, preprocess, pyprocess, scaling, tools, validation # noqa: E402
19from scipy.ndimage import map_coordinates # noqa: E402
22def run_openpiv(
23 image1: str,
24 image2: str,
25 output_dir: str = "results",
26 mask: str = "none",
27 mask_method: str = "intensity",
28 window_size: int = 32,
29 overlap: int = 12,
30 search_area: int = 38,
31 dt: float = 0.02,
32 scaling_factor: float = 96.52,
33 threshold: float = 1.05,
34 drop_invalid: bool = False,
35 verbose: bool = False,
36 backend: str = "scipy",
37) -> Path:
38 """Run PIV analysis on an image pair and return the output directory.
40 scaling_factor is in pixels per physical unit (px/mm for OpenPIV's own test1 data),
41 so u and v come out in that unit per second.
42 """
43 for name, value in (("window_size", window_size), ("search_area", search_area), ("overlap", overlap)):
44 if isinstance(value, bool) or not isinstance(value, (int, np.integer)):
45 raise ValueError(f"{name} must be an integer")
46 if window_size <= 0 or search_area <= 0 or overlap < 0:
47 raise ValueError("window_size/search_area must be positive; overlap must be nonnegative")
48 window_size, search_area, overlap = int(window_size), int(search_area), int(overlap)
49 if search_area < window_size:
50 raise ValueError(
51 f"search_area ({search_area}) must be >= window_size ({window_size})"
52 )
53 if overlap >= search_area:
54 raise ValueError(
55 f"overlap ({overlap}) must be < search_area ({search_area})"
56 )
57 for name, value in (("dt", dt), ("scaling_factor", scaling_factor), ("threshold", threshold)):
58 if not np.isfinite(value) or value <= 0:
59 raise ValueError(f"{name} must be finite and positive")
60 if backend not in ("scipy", "auto", "rust"):
61 raise ValueError("backend must be scipy, auto, or rust")
62 if mask not in ("none", "dynamic"):
63 raise ValueError("mask must be none or dynamic")
64 if mask == "dynamic" and mask_method != "intensity":
65 raise ValueError(
66 "Use mask_method='intensity': OpenPIV 0.26.1 edges masking indexes "
67 "with a uint8 mask and can corrupt the image or raise IndexError."
68 )
70 output_path = Path(output_dir)
71 output_path.mkdir(parents=True, exist_ok=True)
73 if verbose:
74 print(f"Loading images: {image1}, {image2}")
76 frame_a = tools.imread(image1)
77 frame_b = tools.imread(image2)
78 if frame_a.ndim != 2 or frame_b.ndim != 2 or frame_a.shape != frame_b.shape:
79 raise ValueError("Images must be matching 2D grayscale arrays")
80 if min(frame_a.shape) < search_area:
81 raise ValueError("search_area must fit within both image dimensions")
82 if not np.isfinite(frame_a).all() or not np.isfinite(frame_b).all():
83 raise ValueError("Images must contain only finite intensities")
84 image_mask = np.zeros(frame_a.shape, dtype=bool)
86 if mask == "dynamic":
87 if verbose:
88 print(f"Applying dynamic mask (method={mask_method})")
89 # Preserve obstacles separately from correlation-quality flags. They must
90 # remain excluded even if outlier replacement fills a nearby vector.
91 frame_a, mask_a = preprocess.dynamic_masking(
92 frame_a.astype(np.float64), method=mask_method
93 )
94 frame_b, mask_b = preprocess.dynamic_masking(
95 frame_b.astype(np.float64), method=mask_method
96 )
97 image_mask = mask_a.astype(bool) | mask_b.astype(bool)
98 frame_a[image_mask] = 0
99 frame_b[image_mask] = 0
101 u, v, s2n = pyprocess.extended_search_area_piv(
102 frame_a.astype(np.float32),
103 frame_b.astype(np.float32),
104 window_size=window_size,
105 overlap=overlap,
106 dt=dt,
107 search_area_size=search_area,
108 # Linear correlation zero-pads the FFT; normalization follows the
109 # documented linear-correlation recipe. Circular extended search is
110 # also supported upstream, but is not this wrapper's chosen recipe.
111 correlation_method="linear" if search_area > window_size else "circular",
112 normalized_correlation=search_area > window_size,
113 sig2noise_method="peak2peak",
114 backend=backend,
115 )
117 x, y = pyprocess.get_coordinates(
118 image_size=frame_a.shape,
119 search_area_size=search_area,
120 overlap=overlap,
121 center_on_field=False, # match the actual sliding-window extraction
122 )
123 grid_mask = map_coordinates(
124 image_mask.astype(np.uint8), [y, x], order=0, mode="nearest"
125 ).astype(bool)
127 # flags is boolean: True marks a spurious vector.
128 flags = (
129 validation.sig2noise_val(s2n, threshold=threshold)
130 | ~np.isfinite(s2n) | ~np.isfinite(u) | ~np.isfinite(v)
131 )
132 if verbose:
133 total = flags.size
134 print(f"Flagged {int(np.sum(flags))}/{total} vectors below s2n {threshold}")
136 if drop_invalid:
137 u = np.where(flags, np.nan, u)
138 v = np.where(flags, np.nan, v)
139 elif (~(flags | grid_mask)).any():
140 # Solid-body vectors must not contribute to interpolation. Keep their
141 # mask and remove all imputed values inside the body before saving.
142 u = np.ma.array(np.where(grid_mask, np.nan, u), mask=grid_mask)
143 v = np.ma.array(np.where(grid_mask, np.nan, v), mask=grid_mask)
144 u, v = filters.replace_outliers(
145 u, v, flags | grid_mask, method="localmean", max_iter=3, kernel_size=2
146 )
147 u, v = u.filled(np.nan), v.filled(np.nan)
148 else:
149 u, v = np.full(u.shape, np.nan), np.full(v.shape, np.nan)
150 u, v = np.where(grid_mask, np.nan, u), np.where(grid_mask, np.nan, v)
152 x, y, u, v = scaling.uniform(x, y, u, v, scaling_factor=scaling_factor)
153 # Use the image boundary as the physical origin. transform_coordinates only
154 # reverses the grid y values; on an uncentered grid that shifts the overlay.
155 y, v = frame_a.shape[0] / scaling_factor - y, -v
157 np.savez(
158 output_path / "params.npz", x=x, y=y, u=u, v=v, flags=flags,
159 mask=grid_mask, s2n=s2n, dt=dt, scaling_factor=scaling_factor,
160 backend=backend, openpiv_version=version("openpiv"), window_size=window_size,
161 search_area=search_area, overlap=overlap, threshold=threshold,
162 )
164 vectors_file = output_path / "vectors.txt"
165 tools.save(vectors_file, x, y, u, v, flags, mask=grid_mask)
167 fig, ax = plt.subplots(figsize=(8, 8))
168 # The upstream plotting helper estimates image bounds from the last vector,
169 # which stretches overlays when leftover margins or extended windows exist.
170 ax.imshow(
171 tools.imread(image1), cmap="gray", origin="upper",
172 extent=(0, frame_a.shape[1] / scaling_factor, 0, frame_a.shape[0] / scaling_factor),
173 )
174 finite_fluid = ~grid_mask & np.isfinite(u) & np.isfinite(v)
175 for selected, color in ((finite_fluid & ~flags, "b"), (finite_fluid & flags, "r")):
176 if selected.any():
177 ax.quiver(x[selected], y[selected], u[selected], v[selected],
178 color=color, scale=50, width=0.0035)
179 ax.set_aspect("equal")
180 fig.savefig(output_path / "vector_field.png", dpi=150, bbox_inches="tight")
181 plt.close(fig)
183 if verbose:
184 print(f"Results saved to {output_path}")
185 for name in ("vectors.txt", "params.npz", "vector_field.png"):
186 print(f" - {name}")
188 return output_path
191def main():
192 parser = argparse.ArgumentParser(
193 description="OpenPIV - Particle Image Velocimetry processing"
194 )
195 parser.add_argument(
196 "--image",
197 action="append",
198 required=True,
199 help="Image file; specify exactly twice for the pair",
200 )
201 parser.add_argument("--output_dir", default="results", help="Output directory")
202 parser.add_argument(
203 "--mask",
204 default="none",
205 choices=["none", "dynamic"],
206 help="Masking mode (default: none)",
207 )
208 parser.add_argument(
209 "--mask_method",
210 default="intensity",
211 choices=["edges", "intensity"],
212 help="Dynamic mask method; edges is refused due to an upstream 0.26.1 bug",
213 )
214 parser.add_argument(
215 "--window_size", type=int, default=32, help="Window size in pixels"
216 )
217 parser.add_argument("--overlap", type=int, default=12, help="Overlap in pixels")
218 parser.add_argument(
219 "--search_area", type=int, default=38, help="Search area size in pixels"
220 )
221 parser.add_argument(
222 "--dt", type=float, default=0.02, help="Time between frames (s)"
223 )
224 parser.add_argument(
225 "--scaling",
226 type=float,
227 default=96.52,
228 help="Scaling factor, pixels per physical unit (e.g. px/mm)",
229 )
230 parser.add_argument(
231 "--threshold",
232 type=float,
233 default=1.05,
234 help="peak2peak signal-to-noise threshold",
235 )
236 parser.add_argument(
237 "--drop_invalid",
238 action="store_true",
239 help="NaN out flagged vectors instead of keeping interpolated values",
240 )
241 parser.add_argument("--verbose", action="store_true", help="Verbose output")
242 parser.add_argument("--backend", choices=["scipy", "auto", "rust"], default="scipy")
244 args = parser.parse_args()
246 if len(args.image) != 2:
247 parser.error("Exactly two --image arguments required")
249 run_openpiv(
250 image1=args.image[0],
251 image2=args.image[1],
252 output_dir=args.output_dir,
253 mask=args.mask,
254 mask_method=args.mask_method,
255 window_size=args.window_size,
256 overlap=args.overlap,
257 search_area=args.search_area,
258 dt=args.dt,
259 scaling_factor=args.scaling,
260 threshold=args.threshold,
261 drop_invalid=args.drop_invalid,
262 verbose=args.verbose,
263 backend=args.backend,
264 )
267if __name__ == "__main__":
268 main()