1use pyo3::exceptions::{PyRuntimeError, PyValueError};
7use pyo3::prelude::*;
8use pyo3::types::{PyAny, PyDict};
9
10use scirs2_numpy::{IntoPyArray, PyArray2, PyArrayMethods};
12
13use scirs2_ndimage::{
15 analysis::{image_entropy, peak_signal_to_noise_ratio, structural_similarity_index},
17 features::{canny, harris_corners},
19 filters::{
21 bilateral_filter, gaussian_filter, laplace, maximum_filter, median_filter, minimum_filter,
22 sobel, uniform_filter, BorderMode,
23 },
24 interpolation::{rotate, shift, zoom},
26 measurements::{center_of_mass, moments},
28 morphology::{
30 binary_closing, binary_dilation, binary_erosion, binary_opening, distance_transform_edt,
31 grey_dilation, grey_erosion, label,
32 },
33 segmentation::{otsu_threshold, threshold_binary, watershed},
35};
36
37#[pyfunction]
43#[pyo3(signature = (input, sigma, mode="reflect"))]
44fn gaussian_filter_py(
45 py: Python,
46 input: &Bound<'_, PyArray2<f64>>,
47 sigma: f64,
48 mode: &str,
49) -> PyResult<Py<PyArray2<f64>>> {
50 let binding = input.readonly();
51 let data = binding.as_array().to_owned();
52
53 let border_mode = match mode {
54 "reflect" => BorderMode::Reflect,
55 "constant" => BorderMode::Constant,
56 "nearest" => BorderMode::Nearest,
57 "mirror" => BorderMode::Mirror,
58 "wrap" => BorderMode::Wrap,
59 _ => return Err(PyValueError::new_err(format!("Unknown mode: {}", mode))),
60 };
61
62 let result = gaussian_filter(&data, sigma, Some(border_mode), None)
63 .map_err(|e| PyRuntimeError::new_err(format!("Gaussian filter failed: {}", e)))?;
64
65 Ok(result.into_pyarray(py).unbind())
66}
67
68#[pyfunction]
70#[pyo3(signature = (input, size, mode="reflect"))]
71fn median_filter_py(
72 py: Python,
73 input: &Bound<'_, PyArray2<f64>>,
74 size: usize,
75 mode: &str,
76) -> PyResult<Py<PyArray2<f64>>> {
77 let binding = input.readonly();
78 let data = binding.as_array().to_owned();
79
80 let border_mode = match mode {
81 "reflect" => BorderMode::Reflect,
82 "constant" => BorderMode::Constant,
83 "nearest" => BorderMode::Nearest,
84 "mirror" => BorderMode::Mirror,
85 "wrap" => BorderMode::Wrap,
86 _ => return Err(PyValueError::new_err(format!("Unknown mode: {}", mode))),
87 };
88
89 let size_arr = [size, size];
91 let result = median_filter(&data, &size_arr, Some(border_mode))
92 .map_err(|e| PyRuntimeError::new_err(format!("Median filter failed: {}", e)))?;
93
94 Ok(result.into_pyarray(py).unbind())
95}
96
97#[pyfunction]
99#[pyo3(signature = (input, size, mode="reflect"))]
100fn uniform_filter_py(
101 py: Python,
102 input: &Bound<'_, PyArray2<f64>>,
103 size: usize,
104 mode: &str,
105) -> PyResult<Py<PyArray2<f64>>> {
106 let binding = input.readonly();
107 let data = binding.as_array().to_owned();
108
109 let border_mode = match mode {
110 "reflect" => BorderMode::Reflect,
111 "constant" => BorderMode::Constant,
112 "nearest" => BorderMode::Nearest,
113 "mirror" => BorderMode::Mirror,
114 "wrap" => BorderMode::Wrap,
115 _ => return Err(PyValueError::new_err(format!("Unknown mode: {}", mode))),
116 };
117
118 let size_arr = [size, size];
120 let result = uniform_filter(&data, &size_arr, Some(border_mode), None)
121 .map_err(|e| PyRuntimeError::new_err(format!("Uniform filter failed: {}", e)))?;
122
123 Ok(result.into_pyarray(py).unbind())
124}
125
126#[pyfunction]
128#[pyo3(signature = (input, axis=0))]
129fn sobel_py(
130 py: Python,
131 input: &Bound<'_, PyArray2<f64>>,
132 axis: usize,
133) -> PyResult<Py<PyArray2<f64>>> {
134 let binding = input.readonly();
135 let data = binding.as_array().to_owned();
136
137 let result = sobel(&data, axis, None)
138 .map_err(|e| PyRuntimeError::new_err(format!("Sobel filter failed: {}", e)))?;
139
140 Ok(result.into_pyarray(py).unbind())
141}
142
143#[pyfunction]
145fn laplace_py(py: Python, input: &Bound<'_, PyArray2<f64>>) -> PyResult<Py<PyArray2<f64>>> {
146 let binding = input.readonly();
147 let data = binding.as_array().to_owned();
148
149 let result = laplace(&data, None, None)
150 .map_err(|e| PyRuntimeError::new_err(format!("Laplace filter failed: {}", e)))?;
151
152 Ok(result.into_pyarray(py).unbind())
153}
154
155#[pyfunction]
157#[pyo3(signature = (input, sigma_spatial, sigma_intensity))]
158fn bilateral_filter_py(
159 py: Python,
160 input: &Bound<'_, PyArray2<f64>>,
161 sigma_spatial: f64,
162 sigma_intensity: f64,
163) -> PyResult<Py<PyArray2<f64>>> {
164 let binding = input.readonly();
165 let data = binding.as_array().to_owned();
166
167 let result = bilateral_filter(&data, sigma_spatial, sigma_intensity, None)
168 .map_err(|e| PyRuntimeError::new_err(format!("Bilateral filter failed: {}", e)))?;
169
170 Ok(result.into_pyarray(py).unbind())
171}
172
173#[pyfunction]
175#[pyo3(signature = (input, size))]
176fn maximum_filter_py(
177 py: Python,
178 input: &Bound<'_, PyArray2<f64>>,
179 size: usize,
180) -> PyResult<Py<PyArray2<f64>>> {
181 let binding = input.readonly();
182 let data = binding.as_array().to_owned();
183
184 let size_arr = [size, size];
186 let result = maximum_filter(&data, &size_arr, None, None)
187 .map_err(|e| PyRuntimeError::new_err(format!("Maximum filter failed: {}", e)))?;
188
189 Ok(result.into_pyarray(py).unbind())
190}
191
192#[pyfunction]
194#[pyo3(signature = (input, size))]
195fn minimum_filter_py(
196 py: Python,
197 input: &Bound<'_, PyArray2<f64>>,
198 size: usize,
199) -> PyResult<Py<PyArray2<f64>>> {
200 let binding = input.readonly();
201 let data = binding.as_array().to_owned();
202
203 let size_arr = [size, size];
205 let result = minimum_filter(&data, &size_arr, None, None)
206 .map_err(|e| PyRuntimeError::new_err(format!("Minimum filter failed: {}", e)))?;
207
208 Ok(result.into_pyarray(py).unbind())
209}
210
211#[pyfunction]
217#[pyo3(signature = (input, iterations=1))]
218fn binary_erosion_py(
219 py: Python,
220 input: &Bound<'_, PyArray2<u8>>,
221 iterations: usize,
222) -> PyResult<Py<PyArray2<u8>>> {
223 let binding = input.readonly();
224 let data = binding.as_array();
225
226 let bool_data = data.mapv(|x| x != 0);
228
229 let result = binary_erosion(&bool_data, None, Some(iterations), None, None, None, None)
231 .map_err(|e| PyRuntimeError::new_err(format!("Binary erosion failed: {}", e)))?;
232
233 let u8_result = result.mapv(|x| if x { 1u8 } else { 0u8 });
235 Ok(u8_result.into_pyarray(py).unbind())
236}
237
238#[pyfunction]
240#[pyo3(signature = (input, iterations=1))]
241fn binary_dilation_py(
242 py: Python,
243 input: &Bound<'_, PyArray2<u8>>,
244 iterations: usize,
245) -> PyResult<Py<PyArray2<u8>>> {
246 let binding = input.readonly();
247 let data = binding.as_array();
248
249 let bool_data = data.mapv(|x| x != 0);
250
251 let result = binary_dilation(&bool_data, None, Some(iterations), None, None, None, None)
253 .map_err(|e| PyRuntimeError::new_err(format!("Binary dilation failed: {}", e)))?;
254
255 let u8_result = result.mapv(|x| if x { 1u8 } else { 0u8 });
256 Ok(u8_result.into_pyarray(py).unbind())
257}
258
259#[pyfunction]
261#[pyo3(signature = (input, iterations=1))]
262fn binary_opening_py(
263 py: Python,
264 input: &Bound<'_, PyArray2<u8>>,
265 iterations: usize,
266) -> PyResult<Py<PyArray2<u8>>> {
267 let binding = input.readonly();
268 let data = binding.as_array();
269
270 let bool_data = data.mapv(|x| x != 0);
271
272 let result = binary_opening(&bool_data, None, Some(iterations), None, None, None, None)
274 .map_err(|e| PyRuntimeError::new_err(format!("Binary opening failed: {}", e)))?;
275
276 let u8_result = result.mapv(|x| if x { 1u8 } else { 0u8 });
277 Ok(u8_result.into_pyarray(py).unbind())
278}
279
280#[pyfunction]
282#[pyo3(signature = (input, iterations=1))]
283fn binary_closing_py(
284 py: Python,
285 input: &Bound<'_, PyArray2<u8>>,
286 iterations: usize,
287) -> PyResult<Py<PyArray2<u8>>> {
288 let binding = input.readonly();
289 let data = binding.as_array();
290
291 let bool_data = data.mapv(|x| x != 0);
292
293 let result = binary_closing(&bool_data, None, Some(iterations), None, None, None, None)
295 .map_err(|e| PyRuntimeError::new_err(format!("Binary closing failed: {}", e)))?;
296
297 let u8_result = result.mapv(|x| if x { 1u8 } else { 0u8 });
298 Ok(u8_result.into_pyarray(py).unbind())
299}
300
301#[pyfunction]
303#[pyo3(signature = (input, size=3))]
304fn grey_erosion_py(
305 py: Python,
306 input: &Bound<'_, PyArray2<f64>>,
307 size: usize,
308) -> PyResult<Py<PyArray2<f64>>> {
309 let binding = input.readonly();
310 let data = binding.as_array().to_owned();
311
312 let size_arr = [size, size];
314 let result = grey_erosion(&data, Some(&size_arr), None, None, None, None)
315 .map_err(|e| PyRuntimeError::new_err(format!("Grey erosion failed: {}", e)))?;
316
317 Ok(result.into_pyarray(py).unbind())
318}
319
320#[pyfunction]
322#[pyo3(signature = (input, size=3))]
323fn grey_dilation_py(
324 py: Python,
325 input: &Bound<'_, PyArray2<f64>>,
326 size: usize,
327) -> PyResult<Py<PyArray2<f64>>> {
328 let binding = input.readonly();
329 let data = binding.as_array().to_owned();
330
331 let size_arr = [size, size];
333 let result = grey_dilation(&data, Some(&size_arr), None, None, None, None)
334 .map_err(|e| PyRuntimeError::new_err(format!("Grey dilation failed: {}", e)))?;
335
336 Ok(result.into_pyarray(py).unbind())
337}
338
339#[pyfunction]
341fn label_py(py: Python, input: &Bound<'_, PyArray2<u8>>) -> PyResult<Py<PyAny>> {
342 let binding = input.readonly();
343 let data = binding.as_array();
344
345 let bool_data = data.mapv(|x| x != 0);
346 let (labeled, num_features) = label(&bool_data, None, None, None)
348 .map_err(|e| PyRuntimeError::new_err(format!("Label failed: {}", e)))?;
349
350 let dict = PyDict::new(py);
351 dict.set_item("labels", labeled.into_pyarray(py).unbind())?;
352 dict.set_item("num_features", num_features)?;
353
354 Ok(dict.into())
355}
356
357#[pyfunction]
360fn distance_transform_edt_py(
361 py: Python,
362 input: &Bound<'_, PyArray2<u8>>,
363) -> PyResult<Py<PyArray2<f64>>> {
364 let binding = input.readonly();
365 let data = binding.as_array();
366 let shape = data.raw_dim();
367
368 let bool_data = data.mapv(|x| x != 0);
369 let bool_data_dyn = bool_data.into_dyn();
371
372 let (distances_opt, _indices_opt) =
374 distance_transform_edt(&bool_data_dyn, None, true, false)
375 .map_err(|e| PyRuntimeError::new_err(format!("Distance transform failed: {}", e)))?;
376
377 let distances = distances_opt
379 .ok_or_else(|| PyRuntimeError::new_err("Distance transform returned no distances"))?;
380 let result = distances
381 .into_shape_with_order(shape)
382 .map_err(|e| PyRuntimeError::new_err(format!("Failed to reshape result: {}", e)))?;
383
384 Ok(result.into_pyarray(py).unbind())
385}
386
387#[pyfunction]
393#[pyo3(signature = (input, angle, reshape=true, cval=0.0))]
394fn rotate_py(
395 py: Python,
396 input: &Bound<'_, PyArray2<f64>>,
397 angle: f64,
398 reshape: bool,
399 cval: f64,
400) -> PyResult<Py<PyArray2<f64>>> {
401 let binding = input.readonly();
402 let data = binding.as_array().to_owned();
403
404 let result = rotate(
406 &data,
407 angle,
408 None,
409 Some(reshape),
410 None,
411 None,
412 Some(cval),
413 None,
414 )
415 .map_err(|e| PyRuntimeError::new_err(format!("Rotate failed: {}", e)))?;
416
417 Ok(result.into_pyarray(py).unbind())
418}
419
420#[pyfunction]
422#[pyo3(signature = (input, zoom_factor, cval=0.0))]
423fn zoom_py(
424 py: Python,
425 input: &Bound<'_, PyArray2<f64>>,
426 zoom_factor: f64,
427 cval: f64,
428) -> PyResult<Py<PyArray2<f64>>> {
429 let binding = input.readonly();
430 let data = binding.as_array().to_owned();
431
432 let result = zoom(&data, zoom_factor, None, None, Some(cval), None)
434 .map_err(|e| PyRuntimeError::new_err(format!("Zoom failed: {}", e)))?;
435
436 Ok(result.into_pyarray(py).unbind())
437}
438
439#[pyfunction]
441#[pyo3(signature = (input, shift_values, cval=0.0))]
442fn shift_py(
443 py: Python,
444 input: &Bound<'_, PyArray2<f64>>,
445 shift_values: (f64, f64),
446 cval: f64,
447) -> PyResult<Py<PyArray2<f64>>> {
448 let binding = input.readonly();
449 let data = binding.as_array().to_owned();
450
451 let result = shift(
453 &data,
454 &[shift_values.0, shift_values.1],
455 None,
456 None,
457 Some(cval),
458 None,
459 )
460 .map_err(|e| PyRuntimeError::new_err(format!("Shift failed: {}", e)))?;
461
462 Ok(result.into_pyarray(py).unbind())
463}
464
465#[pyfunction]
471fn center_of_mass_py(py: Python, input: &Bound<'_, PyArray2<f64>>) -> PyResult<Py<PyAny>> {
472 let binding = input.readonly();
473 let data = binding.as_array().to_owned();
474
475 let result = center_of_mass(&data)
476 .map_err(|e| PyRuntimeError::new_err(format!("Center of mass failed: {}", e)))?;
477
478 let dict = PyDict::new(py);
479 dict.set_item("center", (result[0], result[1]))?;
480
481 Ok(dict.into())
482}
483
484#[pyfunction]
486#[pyo3(signature = (input, order=3))]
487fn moments_py(py: Python, input: &Bound<'_, PyArray2<f64>>, order: usize) -> PyResult<Py<PyAny>> {
488 let binding = input.readonly();
489 let data = binding.as_array().to_owned();
490
491 let result = moments(&data, order)
493 .map_err(|e| PyRuntimeError::new_err(format!("Moments failed: {}", e)))?;
494
495 let dict = PyDict::new(py);
496 dict.set_item("moments", result.into_pyarray(py).unbind())?;
497
498 Ok(dict.into())
499}
500
501#[pyfunction]
507fn watershed_py(
508 py: Python,
509 input: &Bound<'_, PyArray2<f64>>,
510 markers: &Bound<'_, PyArray2<i32>>,
511) -> PyResult<Py<PyArray2<i32>>> {
512 let input_binding = input.readonly();
513 let input_data = input_binding.as_array().to_owned();
514 let markers_binding = markers.readonly();
515 let markers_data = markers_binding.as_array().to_owned();
516
517 let result = watershed(&input_data, &markers_data)
518 .map_err(|e| PyRuntimeError::new_err(format!("Watershed failed: {}", e)))?;
519
520 Ok(result.into_pyarray(py).unbind())
521}
522
523#[pyfunction]
526#[pyo3(signature = (input, bins=256))]
527fn otsu_threshold_py(
528 py: Python,
529 input: &Bound<'_, PyArray2<f64>>,
530 bins: usize,
531) -> PyResult<Py<PyAny>> {
532 let binding = input.readonly();
533 let data = binding.as_array().to_owned();
534
535 let (binarized, threshold_value) = otsu_threshold(&data, bins)
537 .map_err(|e| PyRuntimeError::new_err(format!("Otsu threshold failed: {}", e)))?;
538
539 let dict = PyDict::new(py);
540 dict.set_item("threshold", threshold_value)?;
541 dict.set_item("binarized", binarized.into_pyarray(py).unbind())?;
542
543 Ok(dict.into())
544}
545
546#[pyfunction]
549fn threshold_binary_py(
550 py: Python,
551 input: &Bound<'_, PyArray2<f64>>,
552 threshold: f64,
553) -> PyResult<Py<PyArray2<f64>>> {
554 let binding = input.readonly();
555 let data = binding.as_array().to_owned();
556
557 let result = threshold_binary(&data, threshold)
559 .map_err(|e| PyRuntimeError::new_err(format!("Threshold failed: {}", e)))?;
560
561 Ok(result.into_pyarray(py).unbind())
562}
563
564#[pyfunction]
571#[pyo3(signature = (input, sigma=1.0, low_threshold=0.1, high_threshold=0.2))]
572fn canny_py(
573 py: Python,
574 input: &Bound<'_, PyArray2<f32>>,
575 sigma: f32,
576 low_threshold: f32,
577 high_threshold: f32,
578) -> PyResult<Py<PyArray2<f32>>> {
579 let binding = input.readonly();
580 let data = binding.as_array().to_owned();
581
582 let result = canny(&data, sigma, low_threshold, high_threshold, None);
584
585 Ok(result.into_pyarray(py).unbind())
586}
587
588#[pyfunction]
591#[pyo3(signature = (input, block_size=2, k=0.04, threshold=0.01))]
592fn harris_corners_py(
593 py: Python,
594 input: &Bound<'_, PyArray2<f32>>,
595 block_size: usize,
596 k: f32,
597 threshold: f32,
598) -> PyResult<Py<PyArray2<u8>>> {
599 let binding = input.readonly();
600 let data = binding.as_array().to_owned();
601
602 let result = harris_corners(&data, block_size, k, threshold);
604
605 let u8_result = result.mapv(|x| if x { 1u8 } else { 0u8 });
607 Ok(u8_result.into_pyarray(py).unbind())
608}
609
610#[pyfunction]
616fn psnr_py(image1: &Bound<'_, PyArray2<f64>>, image2: &Bound<'_, PyArray2<f64>>) -> PyResult<f64> {
617 let img1_binding = image1.readonly();
618 let img2_binding = image2.readonly();
619 let img1_data = img1_binding.as_array();
620 let img2_data = img2_binding.as_array();
621
622 peak_signal_to_noise_ratio(&img1_data, &img2_data)
624 .map_err(|e| PyRuntimeError::new_err(format!("PSNR failed: {}", e)))
625}
626
627#[pyfunction]
629fn ssim_py(image1: &Bound<'_, PyArray2<f64>>, image2: &Bound<'_, PyArray2<f64>>) -> PyResult<f64> {
630 let img1_binding = image1.readonly();
631 let img2_binding = image2.readonly();
632 let img1_data = img1_binding.as_array();
633 let img2_data = img2_binding.as_array();
634
635 structural_similarity_index(&img1_data, &img2_data)
637 .map_err(|e| PyRuntimeError::new_err(format!("SSIM failed: {}", e)))
638}
639
640#[pyfunction]
642fn image_entropy_py(input: &Bound<'_, PyArray2<f64>>) -> PyResult<f64> {
643 let binding = input.readonly();
644 let data = binding.as_array();
645
646 image_entropy(&data)
648 .map_err(|e| PyRuntimeError::new_err(format!("Image entropy failed: {}", e)))
649}
650
651pub fn register_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
653 m.add_function(wrap_pyfunction!(gaussian_filter_py, m)?)?;
655 m.add_function(wrap_pyfunction!(median_filter_py, m)?)?;
656 m.add_function(wrap_pyfunction!(uniform_filter_py, m)?)?;
657 m.add_function(wrap_pyfunction!(sobel_py, m)?)?;
658 m.add_function(wrap_pyfunction!(laplace_py, m)?)?;
659 m.add_function(wrap_pyfunction!(bilateral_filter_py, m)?)?;
660 m.add_function(wrap_pyfunction!(maximum_filter_py, m)?)?;
661 m.add_function(wrap_pyfunction!(minimum_filter_py, m)?)?;
662
663 m.add_function(wrap_pyfunction!(binary_erosion_py, m)?)?;
665 m.add_function(wrap_pyfunction!(binary_dilation_py, m)?)?;
666 m.add_function(wrap_pyfunction!(binary_opening_py, m)?)?;
667 m.add_function(wrap_pyfunction!(binary_closing_py, m)?)?;
668 m.add_function(wrap_pyfunction!(grey_erosion_py, m)?)?;
669 m.add_function(wrap_pyfunction!(grey_dilation_py, m)?)?;
670 m.add_function(wrap_pyfunction!(label_py, m)?)?;
671 m.add_function(wrap_pyfunction!(distance_transform_edt_py, m)?)?;
672
673 m.add_function(wrap_pyfunction!(rotate_py, m)?)?;
675 m.add_function(wrap_pyfunction!(zoom_py, m)?)?;
676 m.add_function(wrap_pyfunction!(shift_py, m)?)?;
677
678 m.add_function(wrap_pyfunction!(center_of_mass_py, m)?)?;
680 m.add_function(wrap_pyfunction!(moments_py, m)?)?;
681
682 m.add_function(wrap_pyfunction!(watershed_py, m)?)?;
684 m.add_function(wrap_pyfunction!(otsu_threshold_py, m)?)?;
685 m.add_function(wrap_pyfunction!(threshold_binary_py, m)?)?;
686
687 m.add_function(wrap_pyfunction!(canny_py, m)?)?;
689 m.add_function(wrap_pyfunction!(harris_corners_py, m)?)?;
690
691 m.add_function(wrap_pyfunction!(psnr_py, m)?)?;
693 m.add_function(wrap_pyfunction!(ssim_py, m)?)?;
694 m.add_function(wrap_pyfunction!(image_entropy_py, m)?)?;
695
696 Ok(())
697}