Add a model selector

app.py

@@ -1,51 +1,59 @@
 # -*- coding: utf-8 -*-
 
 import json
-import os
 from pathlib import Path
 
 import gradio as gr
 import numpy as np
-from doc_ufcn import models
-from doc_ufcn.main import DocUFCN
 from PIL import Image, ImageDraw
 
 from config import parse_configurations
+from tools import UFCNModel
 
 # Load the config
-config = parse_configurations(Path("config.json"))
+config = parse_configurations(Path("config.yaml"))
 
-# Download the model
-model_path, parameters = models.download_model(name=config["model_name"])
-
-# Store classes_colors list
-classes_colors = config["classes_colors"]
+# Check that the paths of the examples are valid
+for example in config["examples"]:
+    assert Path.exists(
+        Path(example)
+    ), f"The path of the image '{example}' does not exist."
+
+# Cached models, maps model_name to UFCNModel object
+MODELS = {
+    model["model_name"]: UFCNModel(
+        name=model["model_name"],
+        colors=model["classes_colors"],
+        title=model["title"],
+        description=model["description"],
+    )
+    for model in config["models"]
+}
 
-# Store classes
-classes = parameters["classes"]
+# Create a list of models name
+models_name = list(MODELS)
 
-# 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."
+def load_model(model_name) -> UFCNModel:
+    """
+    Retrieve the model, and load its parameters/files if it wasn't done before.
 
-# 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"])
+    :param model_name: The name of the selected model
+    :return: The UFCNModel instance selected
+    """
+    assert model_name in MODELS
+    model = MODELS[model_name]
+    # Load the model's files if it wasn't done before
+    if not model.loaded:
+        model.load()
+    return model
 
 
-def query_image(image):
+def query_image(model_name: gr.Dropdown, image: gr.Image) -> list([Image, json]):
     """
-    Draws the predicted polygons with the color provided by the model on an image
+    Loads a model and draws the predicted polygons with the color provided by the model on an image
 
+    :param model: A model selected in dropdown
     :param image: An image to predict
     :return: Image and dict, an image with the predictions and a
         dictionary mapping an object idx (starting from 1) to a dictionary describing the detected object:
@@ -54,8 +62,11 @@ def query_image(image):
         - `channel` key : str, the name of the predicted class.
     """
 
+    # Load the model and get its classes, classes_colors and the model
+    ufcn_model = load_model(model_name)
+
     # Make a prediction with the model
-    detected_polygons, probabilities, mask, overlap = model.predict(
+    detected_polygons, probabilities, mask, overlap = ufcn_model.model.predict(
         input_image=image, raw_output=True, mask_output=True, overlap_output=True
     )
 
@@ -70,12 +81,12 @@ def query_image(image):
 
     # 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 channel in range(1, ufcn_model.num_channels):
         for i, polygon in enumerate(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]
+                polygon["polygon"], fill=ufcn_model.colors[channel - 1]
             )
 
             # Build the dictionary
@@ -88,34 +99,55 @@ def query_image(image):
                     # Confidence that the model predicts the polygon in the right place
                     "confidence": polygon["confidence"],
                     # The channel on which the polygon is predicted
-                    "channel": classes[channel],
+                    "channel": ufcn_model.classes[channel],
                 }
             )
 
     # Return the blend of the images and the dictionary formatted in json
-    return Image.blend(image, img2, 0.5), json.dumps(predict, indent=20)
+    return Image.blend(image, img2, 0.5), json.dumps(predict, indent=2)
 
 
-with gr.Blocks() as process_image:
+def update_model(model_name: gr.Dropdown) -> str:
+    """
+    Update the model title to the title of the current model
 
+    :param model_name: The name of the selected model
+    :return: A new title
+    """
+    return f"## {MODELS[model_name].title}", MODELS[model_name].description
+
+
+with gr.Blocks() as process_image:
     # Create app title
     gr.Markdown(f"# {config['title']}")
 
     # Create app description
     gr.Markdown(config["description"])
 
+    # Create dropdown button
+    model_name = gr.Dropdown(models_name, value=models_name[0], label="Models")
+
+    # get models
+    selected_model: UFCNModel = MODELS[model_name.value]
+
+    # Create model title
+    model_title = gr.Markdown(f"## {selected_model.title}")
+
+    # Create model description
+    model_description = gr.Markdown(selected_model.description)
+
+    # Change model title and description when the model_id is update
+    model_name.change(update_model, model_name, [model_title, model_description])
+
     # Create a first row of blocks
     with gr.Row():
-
         # Create a column on the left
         with gr.Column():
-
             # Generates an image that can be uploaded by a user
             image = gr.Image()
 
             # Create a row under the image
             with gr.Row():
-
                 # Generate a button to clear the inputs and outputs
                 clear_button = gr.Button("Clear", variant="secondary")
 
@@ -124,25 +156,21 @@ with gr.Blocks() as process_image:
 
             # Create a row under the buttons
             with gr.Row():
-
-                # Generate example images that can be used as input image
-                examples = gr.Examples(inputs=image, examples=config["examples"])
+                # Generate example images that can be used as input image for every model
+                gr.Examples(config["examples"], inputs=image)
 
         # Create a column on the right
         with gr.Column():
-
-            # Generates an output image that does not support upload
-            image_output = gr.Image(interactive=False)
+            with gr.Row():
+                # Generates an output image that does not support upload
+                image_output = gr.Image(interactive=False)
 
             # Create a row under the predicted image
             with gr.Row():
-
                 # Create a column so that the JSON output doesn't take the full size of the page
                 with gr.Column():
-
-                    # Create a collapsible region
+                    # # Create a collapsible region
                     with gr.Accordion("JSON"):
-
                         # Generates a json with the model predictions
                         json_output = gr.JSON()
 
@@ -154,7 +182,9 @@ with gr.Blocks() as process_image:
     )
 
     # Create the button to submit the prediction
-    submit_button.click(query_image, inputs=image, outputs=[image_output, json_output])
+    submit_button.click(
+        query_image, inputs=[model_name, image], outputs=[image_output, json_output]
+    )
 
-# Launch the application
+# Launch the application with the public mode (True or False)
 process_image.launch()

config.json

@@ -1,10 +0,0 @@
-{
-    "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

@@ -7,21 +7,22 @@ from teklia_toolbox.config import ConfigParser
 
 def parse_configurations(config_path: Path):
     """
-    Parse multiple JSON configuration files into a single source
+    Parse multiple YAML configuration files into a single source
     of configuration for the HuggingFace app
 
-    :param config_path: pathlib.Path, Path to the .json config file
+    :param config_path: pathlib.Path, Path to the .yaml 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("title")
+    parser.add_option("description")
     parser.add_option("examples", type=list)
+    model_parser = parser.add_subparser("models", many=True)
+
+    model_parser.add_option("model_name")
+    model_parser.add_option("title")
+    model_parser.add_option("description")
+    model_parser.add_option("classes_colors", type=list)
 
     return parser.parse(config_path)

config.yaml

@@ -0,0 +1,30 @@
+---
+title: Teklia - Doc-UFCN Demo
+description: >-
+  [TEKLIA](https://teklia.com/)’s Document Layout Analysis on historical documents. For modern documents, see [ocelus.teklia.com](https://ocelus.teklia.com).
+examples:
+  - resource/hugging_face_1.jpg
+  - resource/hugging_face_2.jpg
+  - resource/hugging_face_3.jpg
+  - resource/hugging_face_4.jpg
+
+models:
+  - model_name: doc-ufcn-generic-historical-line
+    title: Doc-UFCN Generic historical line detection
+    description: >-
+      The [generic historical line detection model](https://huggingface.co/Teklia/doc-ufcn-generic-historical-line) predicts text lines from document images. Please select an image from the examples below or upload your own image!
+    classes_colors:
+      - green
+  - model_name: doc-ufcn-huginmunin-line
+    title: Doc-UFCN Hugin-Munin line detection
+    description: >-
+      The [Hugin-Munin line detection model](https://huggingface.co/Teklia/doc-ufcn-huginmunin-line) predicts horizontal and vertical text lines from Hugin-Munin document images. Please select an image from the examples below or upload your own image!
+    classes_colors:
+      - green
+      - blue
+  - model_name: doc-ufcn-generic-page
+    title: Doc-UFCN Generic page detection
+    description: >-
+      The [generic page detection model](https://huggingface.co/Teklia/doc-ufcn-generic-page) predicts single pages from document images. Please select an image from the examples below or upload your own image!
+    classes_colors:
+      - green

tools.py

@@ -0,0 +1,49 @@
+# -*- coding: utf-8 -*-
+
+from dataclasses import dataclass, field
+
+from doc_ufcn import models
+from doc_ufcn.main import DocUFCN
+
+
+@dataclass
+class UFCNModel:
+    name: str
+    colors: list
+    title: str
+    description: str
+    classes: list = field(default_factory=list)
+    model: DocUFCN = None
+
+    def get_class_name(self, channel_idx):
+        return self.classes[channel_idx]
+
+    @property
+    def loaded(self):
+        return self.model is not None
+
+    @property
+    def num_channels(self):
+        return len(self.classes)
+
+    def load(self):
+        # Download the model
+        model_path, parameters = models.download_model(name=self.name)
+
+        # Store classes
+        self.classes = parameters["classes"]
+
+        # Check that the number of colors is equal to the number of classes -1
+        assert self.num_channels - 1 == len(
+            self.colors
+        ), f"The parameter classes_colors was filled with the wrong number of colors. {self.num_channels-1} colors are expected instead of {len(self.colors)}."
+
+        # Load the model
+        self.model = DocUFCN(
+            no_of_classes=len(self.classes),
+            model_input_size=parameters["input_size"],
+            device="cpu",
+        )
+        self.model.load(
+            model_path=model_path, mean=parameters["mean"], std=parameters["std"]
+        )