From 5102121e0e34b0c7ab386805aeb93fa08ff7762a Mon Sep 17 00:00:00 2001
From: Yohsuke Fukai <ysk@yfukai.net>
Date: Tue, 26 Dec 2023 20:26:23 +0900
Subject: [PATCH] allowed tweaking outlier_predictor

---
 src/m2stitch/_stage_model.py | 17 ++++++++++++++---
 src/m2stitch/stitching.py    | 28 ++++++++++++++++++++++++----
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/src/m2stitch/_stage_model.py b/src/m2stitch/_stage_model.py
index 8a27b23..970b245 100644
--- a/src/m2stitch/_stage_model.py
+++ b/src/m2stitch/_stage_model.py
@@ -1,5 +1,6 @@
 import itertools
 from typing import Callable
+from typing import Optional
 from typing import Tuple
 
 import numpy as np
@@ -11,7 +12,11 @@
 
 
 def compute_image_overlap2(
-    grid: pd.DataFrame, direction: str, sizeY: Int, sizeX: Int, predictor: Callable
+    grid: pd.DataFrame,
+    direction: str,
+    sizeY: Int,
+    sizeX: Int,
+    predictor: Optional[Callable],
 ) -> Tuple[Float, ...]:
     """Compute the value of the image overlap.
 
@@ -25,6 +30,8 @@ def compute_image_overlap2(
         the image width
     sizeX : Int
         the image height
+    predictor : Optional[Callable], optional
+        the predictor function of outliers
 
     Returns
     -------
@@ -43,8 +50,12 @@ def compute_image_overlap2(
         ]
     )
     translation = translation[:, np.all(np.isfinite(translation), axis=0)]
-    c = predictor(translation.T)
-    res = np.median(translation[:, c == 1], axis=1)
+    if predictor is not None:
+        c = predictor(translation.T)
+        ind = c == 1
+    else:
+        ind = np.ones(translation.shape[1], dtype=bool)
+    res = np.median(translation[:, ind], axis=1)
     assert len(res) == 2
     return tuple(res)
 
diff --git a/src/m2stitch/stitching.py b/src/m2stitch/stitching.py
index f1b0a39..5479f3d 100644
--- a/src/m2stitch/stitching.py
+++ b/src/m2stitch/stitching.py
@@ -30,7 +30,7 @@
 
 
 @dataclass
-class ElipticEnvelopPredictor:
+class EllipticEnvelopePredictor:
     contamination: float
     epsilon: float
     random_seed: int
@@ -55,6 +55,7 @@ def stitch_images(
     full_output: bool = False,
     row_col_transpose: bool = True,
     ncc_threshold: Float = 0.5,
+    outlier_predictor_name: str = "elliptic_envelope",
 ) -> Tuple[pd.DataFrame, dict]:
     """Compute image positions for stitching.
 
@@ -95,6 +96,10 @@ def stitch_images(
         the threshold of the normalized cross correlation used to select the initial
         stitched pairs.
 
+    outlier_predictor_name : str, default "elliptic_envelope"
+        the name of the outlier predictor.
+        "elliptic_envelope" or "none" are supported.
+
     Returns
     -------
     grid : pd.DataFrame
@@ -104,6 +109,14 @@ def stitch_images(
     prop_dict : dict
         the dict of estimated parameters. (to be documented)
     """
+
+    if outlier_predictor_name == "elliptic_envelope":
+        outlier_predictor = EllipticEnvelopePredictor(0.1, 0.1, 0)
+    elif outlier_predictor_name == "none":
+        outlier_predictor = None
+    else:
+        raise ValueError(f"Unknown outlier_predictor_name: {outlier_predictor_name}")
+
     images = np.array(images)
     assert (position_indices is not None) or (rows is not None and cols is not None)
     if position_indices is None:
@@ -202,12 +215,19 @@ def get_lims(dimension, size):
     assert np.any(
         grid["left_ncc_first"] > ncc_threshold
     ), "there is no good left pair, (try lowering the ncc_threshold)"
-    predictor = ElipticEnvelopPredictor(contamination=0.4, epsilon=0.01, random_seed=0)
     left_displacement = compute_image_overlap2(
-        grid[grid["left_ncc_first"] > ncc_threshold], "left", sizeY, sizeX, predictor
+        grid[grid["left_ncc_first"] > ncc_threshold],
+        "left",
+        sizeY,
+        sizeX,
+        outlier_predictor,
     )
     top_displacement = compute_image_overlap2(
-        grid[grid["top_ncc_first"] > ncc_threshold], "top", sizeY, sizeX, predictor
+        grid[grid["top_ncc_first"] > ncc_threshold],
+        "top",
+        sizeY,
+        sizeX,
+        outlier_predictor,
     )
     overlap_top = np.clip(100 - top_displacement[0] * 100, pou, 100 - pou)
     overlap_left = np.clip(100 - left_displacement[1] * 100, pou, 100 - pou)