Create app.py

app.py

@@ -0,0 +1,39 @@
+import cv2
+from doc_ufcn.main import DocUFCN
+import gradio as gr
+from doc_ufcn import models
+import numpy as np 
+import random
+from PIL import Image
+
+
+model_path, parameters = models.download_model('generic-page')
+
+model = DocUFCN(len(parameters['classes']), parameters['input_size'], 'cpu')
+model.load(model_path, parameters['mean'], parameters['std'])
+
+def visualize_instance_seg_mask(mask):
+    image = np.zeros((mask.shape[0], mask.shape[1], 3))
+    labels = np.unique(mask)
+    label2color = {label: (random.randint(0, 1), random.randint(0, 255), random.randint(0, 255)) for label in labels}
+    for i in range(image.shape[0]):
+      for j in range(image.shape[1]):
+        image[i, j, :] = label2color[mask[i, j]]
+    image = image / 255
+    return image
+
+def query_image(image):
+    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+    detected_polygons, probabilities, mask, overlap = model.predict(
+    image, raw_output=True, mask_output=True, overlap_output=True
+)
+    return visualize_instance_seg_mask(overlap)
+
+demo = gr.Interface(
+    query_image, 
+    inputs=[gr.Image()], 
+    outputs="image",
+    title="TODO",
+)
+
+demo.launch()