diff --git a/mesonet/data.py b/mesonet/data.py
index 0acdfb7..c7abf85 100644
--- a/mesonet/data.py
+++ b/mesonet/data.py
@@ -6,7 +6,7 @@
 This file has been adapted from data.py in https://github.com/zhixuhao/unet
 """
 from __future__ import print_function
-from keras.preprocessing.image import ImageDataGenerator
+from tensorflow.keras.preprocessing.image import ImageDataGenerator
 import numpy as np
 import os
 import skimage.io as io
diff --git a/mesonet/model.py b/mesonet/model.py
index c1eab81..121a563 100644
--- a/mesonet/model.py
+++ b/mesonet/model.py
@@ -6,8 +6,8 @@
 This file has been adapted from model.py in https://github.com/zhixuhao/unet
 """
 from tensorflow.keras.models import *
-from keras.layers import *
-from keras.optimizers import *
+from tensorflow.keras.layers import *
+from tensorflow.keras.optimizers import *
 
 
 def unet(pretrained_weights=None, input_size=(512, 512, 1)):
diff --git a/mesonet/train_model.py b/mesonet/train_model.py
index 7eade38..0a22e34 100644
--- a/mesonet/train_model.py
+++ b/mesonet/train_model.py
@@ -9,7 +9,7 @@
 from mesonet.data import *
 import numpy as np
 import pandas as pd
-from keras.callbacks import ModelCheckpoint
+from tensorflow.keras.callbacks import ModelCheckpoint
 from mesonet.utils import parse_yaml
 from mesonet.dlc_predict import DLC_edit_bodyparts