Spaces:
Runtime error
Runtime error
import gradio as gr | |
import numpy as np | |
import torch | |
from CCAgT_utils.types.mask import Mask | |
from PIL import Image | |
from torch import nn | |
from transformers import SegformerFeatureExtractor | |
from transformers import SegformerForSemanticSegmentation | |
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') | |
model_hub_name = 'lapix/segformer-b3-finetuned-ccagt-400-300' | |
model = SegformerForSemanticSegmentation.from_pretrained( | |
model_hub_name, | |
).to(device) | |
feature_extractor = SegformerFeatureExtractor.from_pretrained( | |
model_hub_name, | |
) | |
def query_image(image): | |
image = np.array(image) | |
img = Image.fromarray(image) | |
pixel_values = feature_extractor( | |
image, | |
return_tensors='pt', | |
).to(device) | |
with torch.no_grad(): | |
outputs = model(pixel_values=pixel_values) | |
logits = outputs.logits | |
upsampled_logits = nn.functional.interpolate( | |
logits, | |
size=img.size[::-1], # (height, width) | |
mode='bilinear', | |
align_corners=False, | |
) | |
segmentation_mask = upsampled_logits.argmax(dim=1)[0] | |
results = Mask(segmentation_mask).colorized() / 255 | |
return results | |
title = 'SegFormer (b3) - CCAgT dataset' | |
description = f""" | |
This is demo for the SegFormer fine-tuned on sub-dataset from | |
[CCAgT dataset](https://huggingface.co/datasets/lapix/CCAgT). This model | |
was trained to segment cervical cells silver-stained (AgNOR technique) | |
images with resolution of 400x300. The model was available at HF hub at | |
[{model_hub_name}](https://huggingface.co/{model_hub_name}). | |
""" | |
examples = [ | |
[f'https://hf.co/{model_hub_name}/resolve/main/sampleA.png'], | |
[f'https://hf.co/{model_hub_name}/resolve/main/sampleB.png'], | |
] | |
demo = gr.Interface( | |
query_image, | |
inputs=[gr.Image()], | |
outputs='image', | |
title=title, | |
description=description, | |
examples=examples, | |
) | |
demo.launch() | |