import gradio as gr 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-b0-finetuned-ade-512-512" ) model = TFSegformerForSemanticSegmentation.from_pretrained( "nvidia/segformer-b0-finetuned-ade-512-512" ) def ade_palette(): """ADE20K palette that maps each class to RGB values.""" return [ [45, 67, 89], [112, 34, 200], [15, 205, 77], [88, 150, 33], [255, 72, 190], [93, 93, 93], [200, 15, 77], [209, 178, 255], [44, 123, 255], [255, 224, 140], [178, 235, 244], [100, 200, 50], [15, 15, 255], [77, 200, 200], [255, 0, 221], [181, 178, 255], [90, 150, 10], [15, 100, 150], [255, 0, 0], [100, 255, 255], [255, 0, 221], [45, 80, 190], [15, 50, 15], [255, 224, 0], [178, 0, 255], [30, 200, 30], [45, 90, 255], [33, 140, 200], [255, 0, 221], [25, 200, 255], [181, 178, 255], [255, 0, 0], [255, 255, 0], [209, 178, 255], [0, 255, 0], [150, 150, 150], [255, 224, 140], [255, 72, 190], [15, 205, 77], [93, 93, 93], [178, 235, 244], [255, 0, 221], [209, 178, 255], [44, 123, 255], [255, 224, 140], [178, 235, 244], [100, 200, 50], [15, 15, 255], [77, 200, 200], [255, 0, 221], [181, 178, 255], [90, 150, 10], [15, 100, 150], [255, 0, 0], [100, 255, 255], [255, 0, 221], [45, 80, 190], [15, 50, 15], [255, 224, 0], [178, 0, 255], [30, 200, 30], [45, 90, 255], [33, 140, 200], [255, 0, 221], [25, 200, 255], [181, 178, 255], [255, 0, 0], [255, 255, 0], [209, 178, 255], [0, 255, 0], [150, 150, 150], [255, 224, 140], [255, 72, 190], [15, 205, 77], [93, 93, 93], [200, 15, 77], [209, 178, 255], [44, 123, 255], [255, 224, 140], [178, 235, 244], [100, 200, 50], [15, 15, 255], [77, 200, 200], [255, 0, 221], [181, 178, 255], [90, 150, 10], [15, 100, 150], [255, 0, 0], [100, 255, 255], [255, 0, 221], [45, 80, 190], [15, 50, 15], [255, 224, 0], [178, 0, 255], [30, 200, 30], [45, 90, 255], [33, 140, 200], [255, 0, 221], [25, 200, 255], [181, 178, 255], [255, 0, 0], [255, 255, 0], [209, 178, 255], [0, 255, 0], [150, 150, 150], [255, 224, 140], [255, 72, 190], [15, 205, 77], [93, 93, 93], [200, 15, 77], [209, 178, 255], [44, 123, 255], [255, 224, 140], [178, 235, 244], [100, 200, 50], [15, 15, 255], [77, 200, 200], [255, 0, 221], [181, 178, 255], [90, 150, 10], [15, 100, 150], [255, 0, 0], [100, 255, 255], [255, 0, 221], [45, 80, 190], [15, 50, 15], [255, 224, 0], [178, 0, 255], [30, 200, 30], [45, 90, 255], [33, 140, 200], [255, 0, 221], [181, 178, 255], [90, 150, 10], [15, 100, 150], [255, 0, 0], [100, 255, 255], [255, 0, 221], [45, 80, 190], [15, 50, 15], [255, 224, 0], [178, 0, 255], [30, 200, 30], [45, 90, 255], [33, 140, 200], [255, 0, 221], [25, 200, 255], [181, 178, 255], [255, 0, 0], [255, 255, 0], ] 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 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.numpy().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 sepia(input_img): input_img = Image.fromarray(input_img) 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] ) # We reverse the shape of `image` because `image.size` returns width and height. seg = tf.math.argmax(logits, axis=-1)[0] color_seg = np.zeros( (seg.shape[0], seg.shape[1], 3), dtype=np.uint8 ) # height, width, 3 for label, color in enumerate(colormap): color_seg[seg.numpy() == label, :] = color # Show image + mask pred_img = np.array(input_img) * 0.5 + color_seg * 0.5 pred_img = pred_img.astype(np.uint8) fig = draw_plot(pred_img, seg) return fig demo = gr.Interface(fn=sepia, inputs=gr.Image(shape=(400, 600)), outputs=['plot'], examples=["city1.jpg", "city2.jpg", "city3.jpg", "city4.jpg", "city5.jpg"], allow_flagging='never') demo.launch()