Converting from tensorflow frozen graph to tflite for android inference - tensorflow

Question

I'm trying to convert from pytorch to tflite for android inference for a app I'm working on that uses real-time camera data of basketball to create a heatmap of made and missed shots. It's already working for iOS. Here's a demo.
I've managed to convert from pytorch (.pth) to onnx and from onnx to a tensorflow frozen graph (.pb). Inference on that tf frozen graph checks out.
However, when I try to convert from the frozen graph to tflite, I get the following error:
RuntimeError: Inputs and outputs not all float|uint8|int16 types.Node number 2 (ADD) failed to invoke.
Input details from the interpreter [interpreter.get_input_details(), interpreter.get_output_details()] suggest the datatype is numpy.float32, which is where I'm confused. Shouldn't that qualify as float? Any suggestions/help will be much appreciated!
[{'name': 'image', 'index': 21904, 'shape': array([ 3, width, height], dtype=int32), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}]
[{'name': 'action', 'index': 7204, 'shape': array([], dtype=int32), 'dtype': <class 'numpy.float32'>, 'quantization': (0.0, 0)}]
From tensorflow's documentation:
import numpy as np
import tensorflow as tf
# Load TFLite model and allocate tensors.
interpreter = tf.lite.Interpreter(model_path="converted_model.tflite")
interpreter.allocate_tensors()
# Get input and output tensors.
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
# Test model on random input data.
input_shape = input_details[0]['shape']
input_data = np.array(np.random.random_sample(input_shape), dtype=np.float32)
interpreter.set_tensor(input_details[0]['index'], input_data)
interpreter.invoke()
# The function `get_tensor()` returns a copy of the tensor data.
# Use `tensor()` in order to get a pointer to the tensor.
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)
The inference file for tensorflow frozen graph (1.14.0) for anyone in need:
import numpy as np
import tensorflow as tf
from PIL import Image
w = ...
h = ...
class CNN(object):
def __init__(self, model_filepath):
self.model_filepath = model_filepath
self.load_graph(model_filepath = self.model_filepath)
def load_graph(self, model_filepath):
self.graph = tf.Graph()
with tf.gfile.GFile(model_filepath, 'rb') as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
with self.graph.as_default():
self.input = tf.placeholder(np.float32, shape = [3, h, w], name='image')
tf.import_graph_def(graph_def, {'image': self.input})
self.graph.finalize()
self.sess = tf.Session(graph = self.graph)
def test(self, data):
output_tensor = self.graph.get_tensor_by_name('import/action:0')
output = self.sess.run(output_tensor, feed_dict = {self.input: data})
return output
def main():
nn = CNN(model_filepath='out_1.14.pb')
img = np.asarray(Image.open('example.jpg')).astype(np.float32)
img = img.transpose(-1, 0, 1)
ans = nn.test(data=img)
print(ans)
if __name__ == '__main__':
main()