app.py

# -*- coding: utf-8 -*-

import json
from pathlib import Path

import gradio as gr
import numpy as np
from PIL import Image, ImageDraw

from config import parse_configurations
from tools import UFCNModel

# Load the config
config = parse_configurations(Path("config.yaml"))

# 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"]
}

# Create a list of models name
models_name = list(MODELS)


def load_model(model_name) -> UFCNModel:
    """
    Retrieve the model, and load its parameters/files if it wasn't done before.

    :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(model_name: gr.Dropdown, image: gr.Image) -> list([Image, json]):
    """
    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:
        - `polygon` key : list, the coordinates of the points of the polygon,
        - `confidence` key : float, confidence of the model,
        - `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 = ufcn_model.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()

    # Initialize the dictionary which will display the json on the application
    predict = []

    # 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, 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=ufcn_model.colors[channel - 1]
            )

            # Build the dictionary
            # Add an index to dictionary keys to differentiate predictions of the same class
            predict.append(
                {
                    # The list of coordinates of the points of the polygon.
                    # Cast to list of np.int32 to make it JSON-serializable
                    "polygon": np.asarray(polygon["polygon"], dtype=np.int32).tolist(),
                    # Confidence that the model predicts the polygon in the right place
                    "confidence": polygon["confidence"],
                    # The channel on which the polygon is predicted
                    "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=2)


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")

                # Generates a button to submit the prediction
                submit_button = gr.Button("Submit", variant="primary")

            # Create a row under the buttons
            with gr.Row():
                # 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():
            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
                    with gr.Accordion("JSON"):
                        # Generates a json with the model predictions
                        json_output = gr.JSON()

    # Clear button: set default values to inputs and output objects
    clear_button.click(
        lambda: (None, None, None),
        inputs=[],
        outputs=[image, image_output, json_output],
    )

    # Create the button to submit the prediction
    submit_button.click(
        query_image, inputs=[model_name, image], outputs=[image_output, json_output]
    )

# Launch the application with the public mode (True or False)
process_image.launch()