Update app.py

app.py

@@ -65,7 +65,7 @@ def gen_pred(img=inputs, model=seg_model):
     img = tf.expand_dims(img, axis=0)
     pred = model.predict(img)
     pred = np.squeeze(pred, axis=0)
-    return img, pred
+    return pred
 
 
 title = "<h1 style='text-align: center;'>Semantic Segmentation</h1>"
@@ -74,7 +74,7 @@ description = "Upload an image and get prediction mask"
 
 gr.Interface(fn=gen_pred, 
              inputs=inputs, 
-             outputs="image", 
+             outputs=gr.components.Image(), 
              title=title, 
              examples=[["003e2c95d.jpg"], ["003b50a15.jpg"], ["003b48a9e.jpg"], ["0038cbe45.jpg"], ["00371aa92.jpg"]],
              # css=css_code,