diff --git a/.gitignore b/.gitignore
index 779028f..503740a 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,3 +11,4 @@ tests/data/testimages_lowncc.npy
 tests/data/testimages_lowncc2.npy
 tests/data/testimages_lowncc2.csv
 examples/stitched_image.npy
+tests/data/test3/
diff --git a/pyproject.toml b/pyproject.toml
index c987a76..659e9e4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "m2stitch"
-version = "0.5.1"
+version = "0.5.2"
 description = "M2Stitch"
 authors = ["Yohsuke Fukai <ysk@yfukai.net>"]
 license = "MIT"
diff --git a/src/m2stitch/_stage_model.py b/src/m2stitch/_stage_model.py
index ec409f0..bc19ab1 100644
--- a/src/m2stitch/_stage_model.py
+++ b/src/m2stitch/_stage_model.py
@@ -43,7 +43,7 @@ def compute_image_overlap2(
     )
     translation = translation[:, np.all(np.isfinite(translation), axis=0)]
     c = predictor.fit_predict(translation.T)
-    res = np.median(translation[:, c], axis=1)
+    res = np.median(translation[:, c == 1], axis=1)
     assert len(res) == 2
     return tuple(res)