Spaces:
Running
Running
File size: 1,659 Bytes
7ce2ffd 9f889ce 7ce2ffd 9f889ce 7ce2ffd 9f889ce 7ce2ffd 9f889ce 7ce2ffd 9f889ce 7ce2ffd 69785bf 7ce2ffd |
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 58 59 60 61 62 63 64 65 66 67 |
import gradio as gr
import timm
import torch
import torch.nn as nn
class Net2D(nn.Module):
def __init__(self, weights):
super().__init__()
self.backbone = timm.create_model("tf_efficientnet_b6_ns", pretrained=False, global_pool="", num_classes=0)
self.pool_layer = nn.AdaptiveAvgPool2d(1)
self.dropout = nn.Dropout(0.5)
self.linear = nn.Linear(2304, 5)
self.load_state_dict(weights)
def forward(self, x):
x = self.backbone(x)
x = self.pool_layer(x).view(x.size(0), -1)
x = self.dropout(x)
x = self.linear(x)
return x[:, 0] if x.size(1) == 1 else x
def rescale(x):
x = x / 255.0
x = x - 0.5
x = x * 2.0
return x
weights = torch.load("model0.ckpt", map_location=torch.device("cpu"))["state_dict"]
weights = {k.replace("model.", ""): v for k, v in weights.items()}
model = Net2D(weights).eval()
def predict(Image):
img = torch.from_numpy(Image)
img = img.permute(2, 0, 1)
img = img.unsqueeze(0)
img = rescale(img)
with torch.no_grad():
grade = torch.softmax(model(img.float()), dim=1)[0]
cats = ["None", "Mild", "Moderate", "Severe", "Proliferative"]
output_dict = {}
for cat, value in zip(cats, grade):
output_dict[cat] = value.item()
return output_dict
image = gr.Image(shape=(512, 512), image_mode="RGB")
label = gr.Label(label="Grade")
demo = gr.Interface(
fn=predict,
inputs=image,
outputs=label,
examples=["examples/0.png", "examples/1.png", "examples/2.png", "examples/3.png", "examples/4.png"]
)
if __name__ == "__main__":
demo.launch(debug=True)
|