cityscapes / app.py
GoBoKyung
city
5ffae2f
raw
history blame contribute delete
No virus
2.49 kB
import gradio as gr
from PIL import Image
import numpy as np
import tensorflow as tf
from transformers import SegformerFeatureExtractor, TFSegformerForSemanticSegmentation
import os
feature_extractor = SegformerFeatureExtractor.from_pretrained(
"nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
)
model = TFSegformerForSemanticSegmentation.from_pretrained(
"nvidia/segformer-b1-finetuned-cityscapes-1024-1024"
)
def ade_palette():
"""ADE20K palette that maps each class to RGB values."""
return [
[204, 87, 92],
[112, 185, 212],
[45, 189, 106],
[234, 123, 67],
[78, 56, 123],
[210, 32, 89],
[90, 180, 56],
[155, 102, 200],
[33, 147, 176],
[255, 183, 76],
[67, 123, 89],
[190, 60, 45],
[134, 112, 200],
[56, 45, 189],
[200, 56, 123],
[87, 92, 204],
[120, 56, 123],
[45, 78, 123],
[156, 200, 56]
]
labels_list = []
with open(r'labels.txt', 'r') as fp:
for line in fp:
labels_list.append(line[:-1])
colormap = np.asarray(ade_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]
def sepia(input_text):
# Check if the input text is a valid file path
if not os.path.isfile(input_text):
return "Invalid file path. Please enter a valid image file path."
# Load the image using the input text (assumed to be a path to an image)
input_img = Image.open(input_text)
inputs = feature_extractor(images=input_img, return_tensors="tf")
outputs = model(**inputs)
logits = outputs.logits
logits = tf.transpose(logits, [0, 2, 3, 1])
logits = tf.image.resize(
logits, input_img.size[::-1]
)
seg = tf.math.argmax(logits, axis=-1)[0]
color_seg = np.zeros(
(seg.shape[0], seg.shape[1], 3), dtype=np.uint8
)
for label, color in enumerate(colormap):
color_seg[seg.numpy() == label, :] = color
pred_img = np.array(input_img) * 0.5 + color_seg * 0.5
pred_img = pred_img.astype(np.uint8)
# Convert the image array to a Pillow (PIL) image
pred_img = Image.fromarray(pred_img)
return pred_img
# Define the Gradio interface
iface = gr.Interface(fn=sepia, inputs="image", outputs="image")
# Launch the Gradio app
iface.launch()