Spaces:
Running
Running
| # -*- 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() | |