Update app.py

app.py

@@ -54,6 +54,8 @@ seg_model = keras.models.load_model('seg_unet_model.h5', custom_objects={'Combo_
 image_output = gr.outputs.Image(type="pil", label="Output Image")
 # outputs = gr.outputs.HTML() #uncomment for single class output 
 
+rows = 1
+columns = 1
 
 def gen_pred(img, model=seg_model):
     # pil_image = img.convert('RGB')
@@ -67,7 +69,7 @@ def gen_pred(img, model=seg_model):
     pred = model.predict(img)
     pred = np.squeeze(pred, axis=0)
     fig = plt.figure(figsize=(10, 7))
-    fig.add_subplot(rows=1, columns=1, 1)
+    fig.add_subplot(rows, columns, 1)
     plt.imshow(pred, interpolation='catrom')
     plt.axis('off')
     plt.title("Prediction")