lindawangg / COVID-Net

COVID-Net Open Source Initiative
Other
1.15k stars 482 forks source link

Convert the checkpoint model type to TensorFlow Lite or at least saved_model format #162

Closed kapilb7 closed 3 years ago

kapilb7 commented 3 years ago

Hi, can someone help with conversion of the checkpoint model type to ideally, TensorFlow Lite, if not, at least saved_model type. I tried and looked for multiple sources, but kept hitting dead ends...

haydengunraj commented 3 years ago

Hi @kapilb7 ,

This code should work for the CXR-2 model, using TensorFlow 1.15:

import tensorflow as tf

# Model file paths
output_file = 'model_CXR-2.tflite'
meta_file = 'models/COVID-Net_CXR-2/model.meta'
ckpt_path = 'models/COVID-Net_CXR-2/model'

# Model tensor names
# These work for CXR-2, but may need to be changed for other models
image_tensor_name = 'input_1:0'
softmax_tensor_name = 'norm_dense_2/Softmax:0'

graph = tf.Graph()
with graph.as_default():
    sess = tf.Session()
    with sess.as_default():
        # Load graph and pretrained weights
        saver = tf.train.import_meta_graph(meta_file)
        saver.restore(sess, ckpt_path)

        # Get input/output tensors
        image = graph.get_tensor_by_name(image_tensor_name)
        output = graph.get_tensor_by_name(softmax_tensor_name)

        # Convert the model
        converter = tf.lite.TFLiteConverter.from_session(sess, [image], [output])
        tflite_model = converter.convert()

        # Save the model
        with open(output_file, 'wb') as f:
            f.write(tflite_model)

On my end, the resulting model seems to work correctly using this simple test script with a dummy input:

import numpy as np
import tensorflow as tf

# TFLite file path
model_file = 'model_CXR-2.tflite'

# Make model interpreter
interpreter = tf.lite.Interpreter(model_path=model_file)
interpreter.allocate_tensors()

# Get model details
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# Make dummy input for testing purposes
height = input_details[0]['shape'][1]
width = input_details[0]['shape'][2]
input_data = np.ones([1, height, width, 3], dtype=np.float32)

# Run model
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()

# View results
output_data = interpreter.get_tensor(output_details[0]['index']).squeeze()
print(output_data)
kapilb7 commented 3 years ago

Hi, I was able to get a .tflite model from your code, but I'm always getting a positive prediction result when I tried inferencing with a sample image. The Inference.py code which did using the checkpoint model type was able to predict correctly.

import tensorflow as tf
import cv2
import numpy as np

interpreter = tf.lite.Interpreter(model_path="/Users/kapil/Documents/FYP/COVIDNet-CXR/COVIDNet-CXR-2/model_CXR-2.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

class_names = ['negative', 'positive']

image_path='/Users/kapil/Documents/FYP/Normal Lung X-rays (Sample)/76052f7902246ff862f52f5d3cd9cd_jumbo.jpg' 
img = cv2.imread(image_path)
img = cv2.resize(img,(480,480))
img = np.array(img, dtype=np.float32)
img.shape = (1, 480, 480, 3)

interpreter.set_tensor(input_details[0]['index'], img)
interpreter.invoke()

mapping = {'negative': 0, 'positive': 1}
inv_mapping = {0: 'negative', 1: 'positive'}
mapping_keys = list(mapping.keys())

output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)
print('Prediction: {}'.format(inv_mapping[output_data.argmax(axis=1)[0]]))
print('Confidence')
print(' '.join('{}: {:.3f}'.format(cls.capitalize(), output_data[0][i]) for cls, i in mapping.items()))
haydengunraj commented 3 years ago

It looks like you haven't normalized the image to the range [0, 1] that the model expects, which explains why the results differ. Try changing this line:

img = np.array(img, dtype=np.float32)

to:

img = np.array(img, dtype=np.float32)/255.
kapilb7 commented 3 years ago

Cool, that worked. I forgot about that after converting it. 😅

kapilb7 commented 3 years ago

Hi, I tried converting and loading the SavedModel format. I could convert it, but when I try loading the model, I'm getting this error: IndexError: list index (0) out of range

This is the code I ran to convert the checkpoint format to SavedModel:

builder = tf.saved_model.builder.SavedModelBuilder('/Users/kapil/Documents/FYP/COVIDNet-CXR/COVIDNet-CXR-2/savedModel')

builder.add_meta_graph_and_variables(sess, [tf.saved_model.TRAINING, tf.saved_model.SERVING], strip_default_attrs=True)

builder.save()
haydengunraj commented 3 years ago

That's odd, your code seems to work for me. This is the code I ran:

import tensorflow as tf

# Model paths
output_dir = 'models/COVID-Net_CXR-2/savedModel'
meta_file = 'models/COVID-Net_CXR-2/model.meta'
ckpt_path = 'models/COVID-Net_CXR-2/model'

graph = tf.Graph()
with graph.as_default():
    sess = tf.Session()
    with sess.as_default():
        # Load graph and pretrained weights
        saver = tf.train.import_meta_graph(meta_file)
        saver.restore(sess, ckpt_path)

        # Save model
        builder = tf.saved_model.builder.SavedModelBuilder(output_dir)
        builder.add_meta_graph_and_variables(sess, [tf.saved_model.TRAINING, tf.saved_model.SERVING], strip_default_attrs=True)
        builder.save()