File size: 3,130 Bytes
c38a591 8de291e b95fb21 1985ff0 a988bac 8de291e d2d6d64 a988bac 8de291e d2d6d64 8de291e d2d6d64 8de291e d2d6d64 8de291e 955f9f6 4a7ca2e e18d0fd 837aa2c e18d0fd d2d6d64 634ef86 8de291e 8740c76 485c5ce 8740c76 8de291e 634ef86 8de291e a988bac 8de291e 634ef86 8de291e a988bac d2d6d64 8de291e a988bac 8de291e ef7902e 8de291e ef7902e 8de291e 37451d1 8de291e 9b95409 b95fb21 d2d6d64 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
import gradio as gr
from transformers import SegformerFeatureExtractor, SegformerForSemanticSegmentation
import matplotlib.pyplot as plt
from matplotlib import gridspec
import numpy as np
from PIL import Image
import tensorflow as tf
import requests
# Load the pre-trained model and feature extractor
feature_extractor = SegformerFeatureExtractor.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-640-1280")
model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-cityscapes-640-1280")
def my_palette():
return [
[131, 162, 255],
[180, 189, 255],
[255, 227, 187],
[255, 210, 143],
[248, 117, 170],
[255, 223, 223],
[255, 246, 246],
[174, 222, 252],
[150, 194, 145],
[255, 219, 170],
[244, 238, 238],
[50, 38, 83],
[128, 98, 214],
[146, 136, 248],
[255, 210, 215],
[255, 152, 152],
[162, 103, 138],
[63, 29, 56]
]
labels_list = []
with open(r"labels.txt", "r") as fp:
for line in fp:
labels_list.append(line[:-1])
colormap = np.asarray(my_palette())
def greet(input_img):
inputs = feature_extractor(images=input_img, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits
logits_tf = tf.transpose(logits.detach(), [0, 2, 3, 1])
logits_tf = tf.image.resize(
logits_tf, [640, 1280]
)
seg = tf.math.argmax(logits_tf, axis=-1)[0]
color_seg = label_to_color_image(seg.numpy())
# Resize color_seg to match the shape of input_img
color_seg_resized = tf.image.resize(color_seg, (input_img.shape[0], input_img.shape[1]))
pred_img = np.array(input_img) * 0.5 + color_seg_resized * 0.5
# Convert pred_img to NumPy array and then change data type
pred_img = np.array(pred_img).astype(np.uint8)
fig = draw_plot(pred_img, seg.numpy())
return fig
def draw_plot(pred_img, seg):
fig = plt.figure(figsize=(20, 15))
grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])
plt.subplot(grid_spec[0])
plt.imshow(pred_img)
plt.axis("off")
LABEL_NAMES = np.asarray(labels_list)
FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)
FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)
unique_labels = np.unique(seg.astype("uint8"))
ax = plt.subplot(grid_spec[1])
plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")
ax.yaxis.tick_right()
plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])
plt.xticks([], [])
ax.tick_params(width=0.0, labelsize=25)
return fig
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=greet,
inputs="image",
outputs=["plot"],
examples=["image (1).jpg", "image (2).jpg", "image (3).jpg", "image (4).jpg", "image (5).jpg"],
allow_flagging="never"
)
iface.launch(share=True) |