onestone0208/onestone33

1import gradio as gr2from matplotlib import gridspec3import matplotlib.pyplot as plt4import numpy as np5from PIL import Image6import torch7from transformers import AutoImageProcessor, AutoModelForSemanticSegmentation8 9MODEL_ID = "nvidia/segformer-b5-finetuned-cityscapes-1024-1024"10processor = AutoImageProcessor.from_pretrained(MODEL_ID)11model = AutoModelForSemanticSegmentation.from_pretrained(MODEL_ID)12 13def ade_palette():14    """ADE20K palette that maps each class to RGB values."""15    return [16        [0, 0, 0],17        [128, 0, 0],18        [255, 0, 0],19        [0, 128, 0],20        [0, 255, 0],21        [0, 0, 128],22        [0, 0, 255],23        [128, 128, 0],24        [255, 255, 0],25        [128, 0, 128],26        [255, 0, 255],27        [0, 128, 128],28        [0, 255, 255],29        [64, 64, 64],30        [192, 192, 192],31        [255, 128, 0],32        [128, 64, 0],33        [0, 128, 64],34        [72, 128, 64]35    ]36 37labels_list = []38with open("labels.txt", "r", encoding="utf-8") as fp:39    for line in fp:40        labels_list.append(line.rstrip("\n"))41 42colormap = np.asarray(ade_palette(), dtype=np.uint8)43 44def label_to_color_image(label):45    if label.ndim != 2:46        raise ValueError("Expect 2-D input label")47    if np.max(label) >= len(colormap):48        raise ValueError("label value too large.")49    return colormap[label]50 51def draw_plot(pred_img, seg_np):52    fig = plt.figure(figsize=(20, 15))53    grid_spec = gridspec.GridSpec(1, 2, width_ratios=[6, 1])54 55    plt.subplot(grid_spec[0])56    plt.imshow(pred_img)57    plt.axis('off')58 59    LABEL_NAMES = np.asarray(labels_list)60    FULL_LABEL_MAP = np.arange(len(LABEL_NAMES)).reshape(len(LABEL_NAMES), 1)61    FULL_COLOR_MAP = label_to_color_image(FULL_LABEL_MAP)62 63    unique_labels = np.unique(seg_np.astype("uint8"))64    ax = plt.subplot(grid_spec[1])65    plt.imshow(FULL_COLOR_MAP[unique_labels].astype(np.uint8), interpolation="nearest")66    ax.yaxis.tick_right()67    plt.yticks(range(len(unique_labels)), LABEL_NAMES[unique_labels])68    plt.xticks([], [])69    ax.tick_params(width=0.0, labelsize=25)70    return fig71 72def run_inference(input_img):73    # input: numpy array from gradio -> PIL74    img = Image.fromarray(input_img.astype(np.uint8)) if isinstance(input_img, np.ndarray) else input_img75    if img.mode != "RGB":76        img = img.convert("RGB")77 78    inputs = processor(images=img, return_tensors="pt")79    with torch.no_grad():80        outputs = model(**inputs)81        logits = outputs.logits  # (1, C, h/4, w/4)82 83    # resize to original84    upsampled = torch.nn.functional.interpolate(85        logits, size=img.size[::-1], mode="bilinear", align_corners=False86    )87    seg = upsampled.argmax(dim=1)[0].cpu().numpy().astype(np.uint8)  # (H,W)88 89    # colorize & overlay90    color_seg = colormap[seg]                                # (H,W,3)91    pred_img = (np.array(img) * 0.5 + color_seg * 0.5).astype(np.uint8)92 93    fig = draw_plot(pred_img, seg)94    return fig95 96demo = gr.Interface(97    fn=run_inference,98    inputs=gr.Image(type="numpy", label="Input Image"),99    outputs=gr.Plot(label="Overlay + Legend"),100    examples=[101        "images (1).jpeg",102        "images (2).jpeg",103        "images (3).jpeg",104        "image5.jpeg"105    ],106    flagging_mode="never",107    cache_examples=False,108)109 110if __name__ == "__main__":111    demo.launch()112