HEp2 / app.py
admin
2 en
837bdb5
raw
history blame
1.82 kB
import os
import torch
import gradio as gr
from PIL import Image
from torchvision.transforms import transforms
from modelscope import snapshot_download
MODEL_DIR = snapshot_download("Genius-Society/HEp2", cache_dir="./__pycache__")
CLASSES = [
"Centromere",
"Golgi",
"Homogeneous",
"NuMem",
"Nucleolar",
"Speckled",
]
def embeding(img_path: str):
compose = transforms.Compose(
[
transforms.Resize(224),
transforms.CenterCrop(224),
transforms.RandomAffine(5),
transforms.ToTensor(),
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
]
)
img = Image.open(img_path).convert("RGB")
return compose(img)
def infer(target: str):
model = torch.load(f"{MODEL_DIR}/save.pt", map_location=torch.device("cpu"))
if not target:
return None, "Please upload a cell picture!"
torch.cuda.empty_cache()
input: torch.Tensor = embeding(target)
output: torch.Tensor = model(input.unsqueeze(0))
predict = torch.max(output.data, 1)[1]
return os.path.basename(target), CLASSES[predict]
if __name__ == "__main__":
example_imgs = []
for cls in CLASSES:
example_imgs.append(f"{MODEL_DIR}/examples/{cls}.png")
with gr.Blocks() as demo:
gr.Interface(
fn=infer,
inputs=gr.Image(type="filepath", label="Upload a cell picture"),
outputs=[
gr.Textbox(label="Picture name", show_copy_button=True),
gr.Textbox(label="Recognition result", show_copy_button=True),
],
title="It is recommended to upload HEp2 cell images in PNG format.",
examples=example_imgs,
flagging_mode="never",
cache_examples=False,
)
demo.launch()