File size: 3,285 Bytes
e87a462 |
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 |
import numpy as np
import gradio as gr
import cv2
from models.HybridGNet2IGSC import Hybrid
from utils.utils import scipy_to_torch_sparse, genMatrixesLungsHeart
import scipy.sparse as sp
import torch
def getDenseMask(landmarks):
RL = landmarks[0:44]
LL = landmarks[44:94]
H = landmarks[94:]
img = np.zeros([1024,1024], dtype = 'uint8')
RL = RL.reshape(-1, 1, 2).astype('int')
LL = LL.reshape(-1, 1, 2).astype('int')
H = H.reshape(-1, 1, 2).astype('int')
img = cv2.drawContours(img, [RL], -1, 1, -1)
img = cv2.drawContours(img, [LL], -1, 1, -1)
img = cv2.drawContours(img, [H], -1, 2, -1)
return img
def drawOnTop(img, landmarks):
output = getDenseMask(landmarks)
image = np.zeros([1024, 1024, 3])
image[:,:,0] = img + 0.3 * (output == 1).astype('float') - 0.1 * (output == 2).astype('float')
image[:,:,1] = img + 0.3 * (output == 2).astype('float') - 0.1 * (output == 1).astype('float')
image[:,:,2] = img - 0.1 * (output == 1).astype('float') - 0.2 * (output == 2).astype('float')
image = np.clip(image, 0, 1)
return image
def loadModel(device):
A, AD, D, U = genMatrixesLungsHeart()
N1 = A.shape[0]
N2 = AD.shape[0]
A = sp.csc_matrix(A).tocoo()
AD = sp.csc_matrix(AD).tocoo()
D = sp.csc_matrix(D).tocoo()
U = sp.csc_matrix(U).tocoo()
D_ = [D.copy()]
U_ = [U.copy()]
config = {}
config['n_nodes'] = [N1, N1, N1, N2, N2, N2]
A_ = [A.copy(), A.copy(), A.copy(), AD.copy(), AD.copy(), AD.copy()]
A_t, D_t, U_t = ([scipy_to_torch_sparse(x).to(device) for x in X] for X in (A_, D_, U_))
config['latents'] = 64
config['inputsize'] = 1024
f = 32
config['filters'] = [2, f, f, f, f//2, f//2, f//2]
config['skip_features'] = f
hybrid = Hybrid(config.copy(), D_t, U_t, A_t).to(device)
hybrid.load_state_dict(torch.load("weights/bestMSE.pt"))
hybrid.eval()
return hybrid
def pad_to_square(img):
h, w = img.shape[:2]
if h > w:
padw = (h - w)
auxw = padw % 2
img = np.pad(img, ((0, 0), (padw//2, padw//2 + auxw)), 'constant')
padh = 0
auxh = 0
else:
padh = (w - h)
auxh = padh % 2
img = np.pad(img, ((padh//2, padh//2 + auxh), (0, 0)), 'constant')
padw = 0
auxw = 0
return img, (padh, padw, auxh, auxw)
def preprocess(input_img):
img, padding = pad_to_square(input_img)
h, w = img.shape[:2]
if h != 1024 or w != 1024:
img = cv2.resize(img, (1024, 1024), interpolation = cv2.INTER_CUBIC)
return img, (h, w, padding)
def segment(input_img):
input_img = cv2.imread(input_img, 0) / 255.0
img, (h, w, padding) = preprocess(input_img)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
hybrid = loadModel(device)
data = torch.from_numpy(img).unsqueeze(0).unsqueeze(0).to(device).float()
with torch.no_grad():
output = hybrid(data)[0].cpu().numpy().reshape(-1, 2) * 1024
return drawOnTop(img, output)
if __name__ == "__main__":
demo = gr.Interface(segment, gr.Image(type="filepath"), "image")
demo.launch()
|