hugging-face-app (#1)

- Add application files (61a5924e45a0b86be3001966a7c172f6307d7d91)


Co-authored-by: Nolan Boukachab <[email protected]>

app.py

@@ -0,0 +1,86 @@
+# -*- coding: utf-8 -*-
+
+import os
+from pathlib import Path
+import gradio as gr
+from PIL import Image, ImageDraw
+
+from doc_ufcn import models
+from doc_ufcn.main import DocUFCN
+from config import parse_configurations
+
+# Load the config
+config = parse_configurations(Path("config.json"))
+
+# Download the model
+model_path, parameters = models.download_model(name=config["model_name"])
+
+# Store classes_colors list
+classes_colors = config["classes_colors"]
+
+# Store classes
+classes = parameters["classes"]
+
+# Check that the number of colors is equal to the number of classes -1
+assert len(classes) - 1 == len(
+    classes_colors
+), f"The parameter classes_colors was filled with the wrong number of colors. {len(classes)-1} colors are expected instead of {len(classes_colors)}."
+
+# Check that the paths of the examples are valid
+for example in config["examples"]:
+    assert os.path.exists(example), f"The path of the image '{example}' does not exist."
+
+# Load the model
+model = DocUFCN(
+    no_of_classes=len(classes),
+    model_input_size=parameters["input_size"],
+    device="cpu",
+)
+model.load(model_path=model_path, mean=parameters["mean"], std=parameters["std"])
+
+
+def query_image(image):
+    """
+    Draws the predicted polygons with the color provided by the model on an image
+
+    :param image: An image to predict
+    :return: Image, an image with the predictions
+    """
+
+    # Make a prediction with the model
+    detected_polygons, probabilities, mask, overlap = model.predict(
+        input_image=image, raw_output=True, mask_output=True, overlap_output=True
+    )
+
+    # Load image
+    image = Image.fromarray(image)
+
+    # Make a copy of the image to keep the source and also to be able to use Pillow's blend method
+    img2 = image.copy()
+
+    # Create the polygons on the copy of the image for each class with the corresponding color
+    # We do not draw polygons of the background channel (channel 0)
+    for channel in range(1, len(classes)):
+        for polygon in detected_polygons[channel]:
+            # Draw the polygons on the image copy.
+            # Loop through the class_colors list (channel 1 has color 0)
+            ImageDraw.Draw(img2).polygon(
+                polygon["polygon"], fill=classes_colors[channel - 1]
+            )
+
+    # Return the blend of the images
+    return Image.blend(image, img2, 0.5)
+
+
+# Create an interface with the config
+process_image = gr.Interface(
+    fn=query_image,
+    inputs=[gr.Image()],
+    outputs=[gr.Image()],
+    title=config["title"],
+    description=config["description"],
+    examples=config["examples"],
+)
+
+# Launch the application with the public mode (True or False)
+process_image.launch()

config.json

@@ -0,0 +1,10 @@
+{
+    "model_name": "doc-ufcn-generic-historical-line",
+    "classes_colors": ["green"],
+    "title":"doc-ufcn Line Detection Demo",
+    "description":"A demo showing a prediction from the [Teklia/doc-ufcn-generic-historical-line](https://huggingface.co/Teklia/doc-ufcn-generic-historical-line) model. The generic historical line detection model predicts text lines from document images.",
+    "examples":[
+        "resource/hugging_face_1.jpg",
+        "resource/hugging_face_2.jpg"
+    ]
+}

config.py

@@ -0,0 +1,25 @@
+# -*- coding: utf-8 -*-
+
+from pathlib import Path
+from teklia_toolbox.config import ConfigParser
+
+def parse_configurations(config_path: Path):
+    """
+    Parse multiple JSON configuration files into a single source
+    of configuration for the HuggingFace app
+
+    :param config_path: pathlib.Path, Path to the .json config file
+    :return: dict, containing the configuration. Ensures config is complete and with correct typing
+    """
+
+    parser = ConfigParser()
+
+    parser.add_option(
+        "model_name", type=str, default="doc-ufcn-generic-historical-line"
+    )
+    parser.add_option("classes_colors", type=list, default=["green"])
+    parser.add_option("title", type=str)
+    parser.add_option("description", type=str)
+    parser.add_option("examples", type=list)
+
+    return parser.parse(config_path)

requirements.txt

@@ -0,0 +1,2 @@
+doc-ufcn==0.1.9-rc2
+teklia_toolbox==0.1.3