diff --git a/aimodel/src/subcommands/train_predict.py b/aimodel/src/subcommands/train_predict.py
index d65daf7..5bd0e95 100644
--- a/aimodel/src/subcommands/train_predict.py
+++ b/aimodel/src/subcommands/train_predict.py
@@ -11,7 +11,7 @@ from lib.dataset.batched_iterator import batched_iterator
 
 from lib.vis.segmentation_plot import segmentation_plot
 from lib.io.handle_open import handle_open
-from lib.ai.RainfallWaterContraster import RainfallWaterContraster
+from lib.ai.RainfallWaterSegmenter import RainfallWaterSegmenter
 from lib.dataset.dataset_segmenter import dataset_predict
 from lib.io.find_paramsjson import find_paramsjson
 from lib.io.readfile import readfile