| import gradio as gr | |
| from transformers import AutoFeatureExtractor, SegformerForSemanticSegmentation | |
| from PIL import Image | |
| import requests | |
| import matplotlib.pyplot as plt | |
| import torch.nn as nn | |
| extractor = AutoFeatureExtractor.from_pretrained("mattmdjaga/segformer_b2_clothes") | |
| model = SegformerForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes") | |
| def predict(inp): | |
| inputs = extractor(images=inp, return_tensors="pt") | |
| outputs = model(**inputs) | |
| logits = outputs.logits.cpu() | |
| upsampled_logits = nn.functional.interpolate( | |
| logits, | |
| size=inp.size[::-1], | |
| mode="bilinear", | |
| align_corners=False, | |
| ) | |
| pred_seg = upsampled_logits.argmax(dim=1)[0] | |
| pred_seg[pred_seg != 4] = 0 | |
| arr_seg = pred_seg.cpu().numpy().astype("uint8") | |
| arr_seg *= 255 | |
| pil_seg = Image.fromarray(arr_seg) | |
| return pil_seg | |
| gr.Interface(fn=predict, | |
| inputs=gr.Image(type="pil"), | |
| outputs="image", | |
| ).launch() |