|
import gradio as gr |
|
import tensorflow as tf |
|
from PIL import Image |
|
import requests |
|
|
|
|
|
model = tf.saved_model.load("nvidia_segformer_b5_finetuned_cityscapes_1024") |
|
|
|
|
|
label_colors = { |
|
"road": [204, 87, 92], |
|
"sidewalk": [112, 185, 212], |
|
"building": [196, 160, 122], |
|
"wall": [106, 135, 242], |
|
"fence": [91, 192, 222], |
|
"pole": [255, 192, 203], |
|
"traffic_light": [176, 224, 230], |
|
"traffic_sign": [222, 49, 99], |
|
"vegetation": [139, 69, 19], |
|
"terrain": [255, 0, 0], |
|
"sky": [0, 0, 255], |
|
"person": [255, 228, 181], |
|
"rider": [128, 0, 0], |
|
"car": [0, 128, 0], |
|
"truck": [255, 99, 71], |
|
"bus": [0, 255, 0], |
|
"train": [128, 0, 128], |
|
"motorcycle": [255, 255, 0], |
|
"bicycle": [128, 0, 128] |
|
} |
|
|
|
|
|
iface = gr.Interface( |
|
fn=lambda image: predict_segmentation(image, model), |
|
inputs="image", |
|
outputs="image" |
|
) |
|
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() |
|
|