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 [ [111, 214, 93], [18, 181, 57], [72, 152, 135], [240, 74, 253], [211, 22, 184], [68, 111, 215], [120, 212, 135], [185, 244, 20], [190, 90, 92], [53, 18, 220], [251, 56, 67], [141, 248, 248], [226, 38, 196], [153, 75, 248], [158, 166, 127], [240, 254, 73], [157, 99, 218], [85, 243, 54], [38, 71, 123], [207, 188, 66], [145, 24, 6], [187, 252, 239], [240, 181, 229], [137, 187, 112], [104, 219, 158], [234, 56, 176], [23, 141, 13], [28, 22, 88], [83, 169, 127], [1, 236, 221], [61, 88, 81], [102, 94, 10], [116, 233, 66], [147, 247, 143], [241, 72, 39], [229, 165, 195], [22, 247, 217], [110, 208, 164], [236, 236, 6], [163, 31, 15], [78, 148, 190], [92, 222, 66], [198, 120, 99], [161, 201, 28], [235, 88, 53], [249, 233, 102], [235, 115, 89], [51, 135, 171], [37, 162, 46], [11, 200, 171], [192, 186, 65], [173, 208, 139], [240, 124, 1], [106, 209, 96], [174, 126, 239], [221, 234, 164], [140, 46, 109], [135, 62, 174], [130, 51, 242], [229, 28, 133], [30, 157, 217], [154, 195, 123], [157, 115, 35], [199, 218, 59], [144, 47, 157], [253, 185, 226], [8, 62, 238], [71, 191, 146], [217, 227, 170], [169, 195, 73], [253, 60, 179], [42, 239, 174], [67, 221, 248], [163, 179, 218], [250, 30, 153], [154, 66, 181], [109, 228, 192], [213, 212, 73], [125, 186, 185], [12, 80, 88], [188, 90, 227], [38, 131, 95], [105, 56, 175], [230, 72, 244], [212, 98, 68], [5, 14, 131], [136, 150, 164], [72, 70, 198], [160, 124, 189], [255, 132, 160], [199, 71, 86], [32, 209, 66], [167, 50, 228], [163, 72, 61], [53, 24, 145], [132, 27, 124], [72, 143, 166], [54, 156, 177], [197, 26, 37], [230, 92, 201], [31, 47, 165], [133, 215, 89], [190, 51, 145], [162, 3, 41], [37, 197, 236], [247, 19, 29], [105, 12, 99], [130, 235, 57], [112, 224, 59], [6, 253, 14], [205, 176, 152], [110, 202, 51], [94, 74, 61], [108, 86, 56], [148, 184, 162], [125, 0, 195], [143, 211, 60], [108, 240, 95], [106, 211, 59], [12, 1, 158], [46, 53, 36], [130, 192, 113], [204, 224, 85], [162, 86, 98], [10, 155, 230], [76, 105, 166], [157, 34, 206], [3, 230, 115], [115, 172, 117], [98, 2, 191], [173, 132, 102], [3, 47, 51], [60, 7, 102], [70, 47, 237], [10, 145, 167], [235, 156, 244], [142, 188, 86], [137, 45, 182], [110, 37, 249], [21, 108, 156], [51, 19, 187], [66, 99, 230], [249, 153, 221], [231, 146, 194], [153, 115, 50], [25, 15, 226], [126, 9, 119], [241, 114, 28], [134, 156, 64], [111, 215, 120], ] 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] unique_labels = np.asarray([]) def draw_plot(pred_img, seg): global unique_labels 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): global unique_labels 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) outputStr = f"이번에는 " for i in unique_labels: outputStr += labels_list[i] + ", " outputStr += "가 검출됐어요." return fig, outputStr demo = gr.Interface( fn=sepia, inputs=gr.Image(shape=(800, 600)), outputs=["plot", "text"], examples=[ "image (1).jpg", "image (2).jpg", "image (3).jpg", "image (4).jpg", "image (5).jpg"], allow_flagging="never", ) demo.launch()