Spaces:
Sleeping
Sleeping
File size: 1,749 Bytes
518a32b |
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 |
import gradio as gr
import torch
from transformers import SegformerForSemanticSegmentation, SegformerImageProcessor
import numpy as np
from PIL import Image
# Load model and processor
model_name = "nvidia/segformer-b0-finetuned-ade-512-512"
processor = SegformerImageProcessor.from_pretrained(model_name)
model = SegformerForSemanticSegmentation.from_pretrained(model_name)
# ADE20k color palette (simplified for visualization)
def create_color_map():
# Create a colorful palette for visualization
np.random.seed(42)
return np.random.randint(0, 255, size=(150, 3), dtype=np.uint8)
color_map = create_color_map()
def segment_image(image):
# Process the image
inputs = processor(images=image, return_tensors="pt")
# Get model prediction
with torch.no_grad():
outputs = model(**inputs)
logits = outputs.logits
# Get segmentation map
seg = logits.argmax(dim=1)[0].cpu().numpy()
# Convert to colored segmentation map
colored_seg = np.zeros((seg.shape[0], seg.shape[1], 3), dtype=np.uint8)
for label in np.unique(seg):
colored_seg[seg == label] = color_map[label]
# Convert to PIL image for display
segmented_image = Image.fromarray(colored_seg)
return [image, segmented_image]
# Create Gradio interface
demo = gr.Interface(
fn=segment_image,
inputs=gr.Image(type="pil"),
outputs=[
gr.Image(type="pil", label="Original"),
gr.Image(type="pil", label="Segmented")
],
title="Image Segmentation with SegFormer",
description="Upload an image to segment it into different semantic regions using SegFormer model fine-tuned on ADE20K dataset."
)
# Launch the app
if __name__ == "__main__":
demo.launch()
|