|
import gradio as gr |
|
import random |
|
|
|
from matplotlib import gridspec |
|
import matplotlib.pyplot as plt |
|
import numpy as np |
|
from PIL import Image |
|
import tensorflow as tf |
|
from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation |
|
|
|
feature_extractor = SegformerFeatureExtractor.from_pretrained( |
|
"nvidia/segformer-b5-finetuned-cityscapes-1024-1024" |
|
) |
|
model = TFSegformerForSemanticSegmentation.from_pretrained( |
|
"nvidia/segformer-b5-finetuned-cityscapes-1024-1024" |
|
) |
|
|
|
def palette(): |
|
|
|
return [ |
|
[204, 87, 92], |
|
[112, 185, 212], |
|
[196, 160, 122], |
|
[106, 135, 242], |
|
[91, 192, 222], |
|
[255, 192, 203], |
|
[176, 224, 230], |
|
[222, 49, 99], |
|
[139, 69, 19], |
|
[255, 0, 0], |
|
[0, 0, 255], |
|
[255, 228, 181], |
|
[128, 0, 0], |
|
[0, 128, 0], |
|
[255, 99, 71], |
|
[0, 255, 0], |
|
[128, 0, 128], |
|
[255, 255, 0], |
|
[128, 0, 128] |
|
|
|
] |
|
|
|
labels_list = [] |
|
|
|
with open(r'labels.txt', 'r') as fp: |
|
for line in fp: |
|
labels_list.append(line[:-1]) |
|
colormap = np.asarray(palette()) |
|
|
|
def label_to_color_image(label): |
|
if label.ndim != 2: |
|
raise ValueError("Expect 2-D input label") |
|
|
|
if np.max(label) >= len(colormap): |
|
raise ValueError("label value too large.") |
|
return colormap[label] |
|
|
|
|
|
iface = gr.Interface( |
|
fn=predict_segmentation, |
|
inputs=gr.Image(shape=(564,846)), |
|
outputs="image", |
|
examples=["city1.jpg","city2.jpg","city3.jpg"], |
|
) |
|
iface.launch() |
|
|
|
|
|
def predict_segmentation(image, model): |
|
|
|
image = Image.fromarray(image.astype('uint8'), 'RGB') |
|
image = image.resize((1024, 1024)) |
|
image_array = tf.keras.preprocessing.image.img_to_array(image) |
|
image_array = tf.expand_dims(image_array, 0) |
|
|
|
|
|
predictions = model(image_array)["output_0"] |
|
|
|
|
|
segmented_image = tf.zeros_like(predictions) |
|
for label, color in label_colors.items(): |
|
mask = tf.reduce_all(tf.equal(predictions, color), axis=-1, keepdims=True) |
|
for i in range(3): |
|
segmented_image += tf.cast(mask, tf.float32) * tf.constant(color[i], dtype=tf.float32) |
|
|
|
|
|
segmented_image = tf.cast(segmented_image, tf.uint8) |
|
segmented_image = tf.image.resize(segmented_image, [image.height, image.width]) |
|
return segmented_image.numpy() |
|
|