Spaces:
Runtime error
Runtime error
File size: 4,949 Bytes
01194aa 6507855 01194aa 7b5a387 01194aa 7b5a387 01194aa 7b5a387 01194aa 98275ef 01194aa e2eb5fd 6e159e4 e2eb5fd 01194aa 6b4f54d 01194aa b2f4127 01194aa eb2d925 |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 |
import os
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
from torchvision import models
from torchvision import transforms
from PIL import Image
from PIL import UnidentifiedImageError
import cv2
import segmentation_models_pytorch as sm
import gradio as gr
from aib_cdr import calculate_cdr
class Net1(nn.Module):
def __init__(self):
super(Net1, self).__init__()
self.model = models.resnet50(pretrained=True)
for param in self.model.parameters():
param.requires_grad = False
num_ftrs = self.model.fc.in_features
self.model.fc = nn.Sequential(
nn.Linear(num_ftrs, 512),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(512, 3),
nn.Softmax(dim=1)
)
def forward(self, x):
x = self.model(x)
return x
class Net2(nn.Module):
def __init__(self):
super(Net2, self).__init__()
self.model = models.resnet50(pretrained=True)
for param in self.model.parameters():
param.requires_grad = False
num_ftrs = self.model.fc.in_features
self.model.fc = nn.Sequential(
nn.Linear(num_ftrs, 512),
nn.ReLU(),
nn.Dropout(0.2),
nn.Linear(512, 2)
)
def forward(self, x):
x = self.model(x)
probas = nn.functional.softmax(x, dim=1)
return x, probas
disc_model_path = 'models/disc_model.pth'
cup_model_path = 'models/cup_model.pth'
quality_model_path = 'models/quality_model.pth'
camtype_model_path = 'models/camtype_model.pth'
dr_model_path = 'models/dr_model.pth'
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
disc_model = sm.Unet('resnet34', classes=2, activation=None)
disc_model.load_state_dict(torch.load(disc_model_path, map_location=device))
disc_model.to(device)
disc_model.eval()
cup_model = sm.Unet('resnet34', classes=2, activation=None)
cup_model.load_state_dict(torch.load(cup_model_path, map_location=device))
cup_model.to(device)
cup_model.eval()
quality_model = Net1().to(device)
quality_model.load_state_dict(torch.load(quality_model_path,map_location=device))
quality_model.to(device)
quality_model.eval()
camtype_model = models.resnet50(pretrained=True)
num_features = camtype_model.fc.in_features
camtype_model.fc = nn.Linear(num_features, 1)
camtype_model.load_state_dict(torch.load(camtype_model_path,map_location=device))
camtype_model.to(device)
camtype_model.eval()
dr_model = Net2().to(device)
dr_model.load_state_dict(torch.load(dr_model_path,map_location=device))
dr_model.to(device)
dr_model.eval()
transform = transforms.Compose([
transforms.Resize((512, 512)),
transforms.ToTensor(),
])
def model_interface(image):
image = Image.fromarray(image.astype('uint8'), 'RGB')
image = image.resize((512, 512))
image_np = np.array(image)
image_tensor = transform(image).unsqueeze(0).to(device)
cdr = calculate_cdr(image_tensor, disc_model, cup_model, device)
cdr = round(cdr,2)
quality_probs = quality_model(image_tensor)
quality_probs = quality_probs.flatten()
quality_dict = {"Image Quality : ACCEPTABLE": quality_probs[0].item(), "Image Quality : GOOD": quality_probs[1].item(),"Image Quality : POOR" : quality_probs[2].item()}
camtype_pred = camtype_model(image_tensor)
camtype_pred = camtype_pred.item()
camtype = "Eidon" if camtype_pred == 0 else "Nidek"
dr_grading = dr_model(image_tensor)
dr_grading = dr_grading[1].flatten()
dr_dict = {"DR : Negative": dr_grading[0].item(), "DR : Positive": dr_grading[1].item()}
disc_pred = disc_model(image_tensor)
disc_mask = (disc_pred.argmax(dim=1) > 0.5).squeeze().cpu().numpy()
disc_contours, _ = cv2.findContours(disc_mask.astype('uint8'), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
image_np1 = np.array(image)
image_np1 = cv2.drawContours(image_np1, disc_contours, -1, (0, 255, 0), 2)
cup_pred = cup_model(image_tensor)
cup_mask = (cup_pred.argmax(dim=1) > 0.5).squeeze().cpu().numpy()
cup_contours, _ = cv2.findContours(cup_mask.astype('uint8'), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
image_np2 = np.array(image)
image_np2 = cv2.drawContours(image_np2, cup_contours, -1, (0, 0, 255), 2)
return quality_dict, camtype, dr_dict, image_np1 , image_np2,cdr
iface = gr.Interface(
fn=model_interface,
inputs=gr.inputs.Image(),
outputs=[
gr.outputs.Label(num_top_classes=3, label="Image Quality"),
gr.outputs.Textbox(label="Camera Type"),
gr.outputs.Label(num_top_classes=2, label="DR Grading"),
gr.outputs.Image(type='numpy', label="Disc Mask"),
gr.outputs.Image(type='numpy', label="Cup Mask"),
gr.outputs.Textbox(label="Cup-to-disc Ratio")
]
)
iface.launch(debug=True) |