new model working

app.py

@@ -6,12 +6,15 @@ from itertools import islice
 from PIL import Image, ImageDraw, ImageColor
 from line_fit import LineFit
 import random
+import lwbna_unet as u
 
 MAX_SELECTIONS = 8
 input_shape = (MAX_SELECTIONS,256,256,2)
 def load_model():
-    model = tf.keras.models.load_model('vivid-sweep-model-best.hdf5')
-    return model
+    unet = u.LWBNAUnet(1,64,4,4)
+    unet.build(input_shape)
+    unet.load_weights('robust-energy-model-best.hdf5')
+    return unet
 
 model = load_model()
 colors = list(ImageColor.colormap.keys())

requirements.txt

@@ -1,4 +1,5 @@
 tensorflow
 opencv-python
 gradio
-scipy
+scipy
+lwbna-unet