33 lines
1.3 KiB
Python
33 lines
1.3 KiB
Python
import tensorflow as tf
|
|
import os
|
|
import numpy as np
|
|
from .processImg import process_image_file
|
|
def detectxray(imagepath):
|
|
weightspath = "models/COVIDNet-CXR4-A"
|
|
metaname = "model.meta"
|
|
ckptname = "model-18540"
|
|
n_classes = "3"
|
|
in_tensorname = "input_1:0"
|
|
out_tensorname = "norm_dense_1/Softmax:0"
|
|
input_size = 480
|
|
top_percent = 0.08
|
|
mapping = {'normal': 0, 'pneumonia': 1, 'COVID-19': 2}
|
|
inv_mapping = {0: 'normal', 1: 'pneumonia', 2: 'COVID-19'}
|
|
mapping_keys = list(mapping.keys())
|
|
sess = tf.Session()
|
|
tf.get_default_graph()
|
|
saver = tf.train.import_meta_graph(os.path.join(weightspath, metaname))
|
|
saver.restore(sess, os.path.join(weightspath, ckptname))
|
|
graph = tf.get_default_graph()
|
|
image_tensor = graph.get_tensor_by_name(in_tensorname)
|
|
pred_tensor = graph.get_tensor_by_name(out_tensorname)
|
|
x = process_image_file(imagepath, input_size, top_percent=top_percent)
|
|
x = x.astype('float32') / 255.0
|
|
feed_dict = {image_tensor: np.expand_dims(x, axis=0)}
|
|
pred = sess.run(pred_tensor, feed_dict=feed_dict)
|
|
prediction = inv_mapping[pred.argmax(axis=1)[0]]
|
|
pred_normal = round(pred[0][mapping['normal']], 3)
|
|
pred_pneu = round(pred[0][mapping['pneumonia']], 3)
|
|
pred_covid = round(pred[0][mapping['COVID-19']], 3)
|
|
return prediction,pred_normal,pred_pneu,pred_covid
|