Update app.py

app.py

@@ -1,3 +1,4 @@
+import os
 import gradio as gr
 import numpy as np
 import tensorflow as tf
@@ -5,74 +6,67 @@ from skimage.io import imsave
 from skimage.transform import resize
 import matplotlib.pyplot as plt
 
-#from SegZoo
-def standardize(img):
-    #standardization using adjusted standard deviation
+# Suppress TensorFlow warnings
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 
+# Standardize function
+def standardize(img):
+    # Standardization using adjusted standard deviation
     N = np.shape(img)[0] * np.shape(img)[1]
-    s = np.maximum(np.std(img), 1.0/np.sqrt(N))
+    s = np.maximum(np.std(img), 1.0 / np.sqrt(N))
     m = np.mean(img)
     img = (img - m) / s
-    del m, s, N
-    #
-    if np.ndim(img)==2:
-        img = np.dstack((img,img,img))
+    if np.ndim(img) == 2:
+        img = np.dstack((img, img, img))
     return img
 
-#load model
+# Load model
 filepath = './saved_model'
-model = tf.keras.models.load_model(filepath, compile = True)
-model.compile
+model = tf.keras.layers.TFSMLayer(filepath, call_endpoint='serving_default')
 
-#segmentation
+# Segmentation function
 def FRFsegment(input_img):
-    
-    dims=(512,512)
-    w = input_img.shape[0]
-    h = input_img.shape[1]  
-    print(w)
-    print(h)
+    dims = (512, 512)
+    w, h = input_img.shape[:2]
 
     img = standardize(input_img)
-    img = resize(img, dims, preserve_range=True, clip=True) 
-    img = np.expand_dims(img,axis=0)
-    
-    est_label = model.predict(img)
-    
-#     # Test Time AUgmentation   
-#     est_label2 = np.flipud(model.predict((np.flipud(img)), batch_size=1))
-#     est_label3 = np.fliplr(model.predict((np.fliplr(img)), batch_size=1))
-#     est_label4 = np.flipud(np.fliplr(model.predict((np.flipud(np.fliplr(img))))))
+    img = resize(img, dims, preserve_range=True, clip=True)
+    img = np.expand_dims(img, axis=0)
+
+    est_label = model(img)
 
-#     #soft voting - sum the softmax scores to return the new TTA estimated softmax scores
-#     pred = est_label + est_label2 + est_label3 + est_label4
-#     est_label = pred
-    
-    mask = np.argmax(np.squeeze(est_label, axis=0),-1)
+    # Mask creation
+    mask = np.argmax(np.squeeze(est_label, axis=0), -1)
     pred = resize(mask, (w, h), preserve_range=True, clip=True)
 
     # Convert pred to uint8
     pred_uint8 = (pred / np.max(pred) * 255).astype(np.uint8)
 
     imsave("label.png", pred_uint8)
-    
-    #overlay plot
+
+    # Overlay plot
     plt.clf()
-    plt.imshow(input_img,cmap='gray')
+    plt.imshow(input_img, cmap='gray')
     plt.imshow(pred, alpha=0.4)
     plt.axis("off")
     plt.margins(x=0, y=0)
     plt.savefig("overlay.png", dpi=300, bbox_inches="tight")
-    
-    return plt, "label.png", "overlay.png"
-
-out1 = gr.outputs.File()
-out2 = gr.outputs.File()
 
+    return "label.png", "overlay.png"
 
+# Gradio Interface
 title = "Segment beach imagery taken from a tower in Duck, NC, USA"
 description = "This model segments beach imagery into 4 classes: vegetation, sand, coarse sand, and background (water + sky + buildings + people)"
-examples = [['examples/FRF_c1_snap_20191112160000.jpg'],['examples/FRF_c1_snap_20170101.jpg']] 
+examples = [['examples/FRF_c1_snap_20191112160000.jpg'], ['examples/FRF_c1_snap_20170101.jpg']]
+
+FRFSegapp = gr.Interface(
+    fn=FRFsegment, 
+    inputs=gr.Image(type="numpy"), 
+    outputs=[gr.File(), gr.File()], 
+    examples=examples, 
+    title=title, 
+    description=description
+)
 
+FRFSegapp.launch()
 
-FRFSegapp = gr.Interface(FRFsegment, gr.inputs.Image(), ['plot',out1, out2], examples=examples, title = title, description = description).launch()