Update app.py

app.py

@@ -14,21 +14,14 @@ print('Successfully loaded model...')
 examples = ['examples/VinDr_RibCXR_val_008.png', 'examples/VinDr_RibCXR_val_013.png']
 
 
-def infer(image):
-    inputs = load_images([image])[..., 1:2]
-    logits = predict(model, inputs)
-    plasma = plt.get_cmap('plasma')
+def segment(img_arr):
+    if np.ndim(img_arr) != 2:
+        img_arr = img_arr[:, :, 1]
+    image = img_arr[np.newaxis, :, :, np.newaxis]
     threshold = 0.5
+    logits = model(image)
     prob = tf.sigmoid(logits)
     pred = tf.cast(prob > threshold, dtype=tf.float32)
-    pred = pred.numpy())[0,:,:,0]
-    image_out = plasma(rescaled)[:, :, :3]
-    return image_out
-
-iface = gr.Interface(
-    fn=infer,
-    title="ribs segmentation",
-    description = "Keras Implementation of Unet++ architecture for ribs segmentation 📏",
-    inputs=[gr.inputs.Image(label="image", type="numpy", shape=(640, 480))],
-    outputs="image",
-    examples=examples).launch(debug=True, cache_examples=True)
+    pred = np.array(pred.numpy())[0,:,:,0]
+    return pred
+iface = gr.Interface(fn=segment, inputs=gr.inputs.Image(shape=(512, 512)), outputs="image", flagging_dir=os.path.join(sys.path[0], "flagged")).launch(share=True)