draft

app.py

@@ -2,36 +2,37 @@ 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
+from PIL import ImageDraw
 
+model_path, parameters = models.download_model("generic-page")
 
-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"])
 
-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)
+        image, raw_output=True, mask_output=True, overlap_output=True
+    )
+    max_confidence = 0
+    max_prediction = None
+    for prediction in detected_polygons[1]:
+        if prediction["confidence"] > max_confidence:
+            max_confidence = prediction["confidence"]
+            max_prediction = prediction
+    image = Image.fromarray(image)
+    img2 = image.copy()
+    xy = max_prediction["polygon"]
+    draw = ImageDraw.Draw(img2)
+    draw.polygon(xy, fill="blue")
+    return Image.blend(image, img2, 0.5)
+
 
 demo = gr.Interface(
-    query_image, 
-    inputs=[gr.Image()], 
+    query_image,
+    inputs=[gr.Image()],
     outputs="image",
     title="TODO",
 )