Support TFLite inference

Author: RyanDam

fdfat/engine/base.py

@@ -90,4 +90,5 @@ def load_checkpoint(self, checkpoint, map_location='cpu', epoch_info=True):
             self.start_epoch = checkpoint['epoch']
             self.best_epoch_loss = checkpoint['best_fit']
             self.best_epoch_no = checkpoint['best_epoch']
-        LOGGER.info(f"Loaded checkpoint epoch {checkpoint['epoch']}")
+        LOGGER.info(f"Loaded checkpoint epoch {checkpoint['epoch']}")
+        

fdfat/nn/onnx.py fdfat/nn/infer.py

@@ -1,45 +1,84 @@
 import cv2
 import numpy as np
-import onnxruntime as ort
 
 from fdfat.utils import box_utils
 
 # ONNX_BACKENDS = ['CoreMLExecutionProvider', 'CPUExecutionProvider']
 # ONNX_BACKENDS = ['CoreMLExecutionProvider']
 ONNX_BACKENDS = ['CPUExecutionProvider']
 
-class ONNXModel:
+class InferModelBackend:
 
-    def __init__(self, model_path, channel_first=True):
+    ONNX = 0
+    TFLITE = 1
+
+class InferModel:
+
+    def __init__(self, model_path, channel_first=True, backend=InferModelBackend.ONNX):
 
         self.model_path = model_path
-        # self.input_width, self.input_height = input_size
+        self.backend = backend
+        self.channel_first = channel_first
+
+        if self.backend == InferModelBackend.ONNX:
+
+            import onnxruntime as ort
+
+            sess_options = ort.SessionOptions()
+            sess_options.intra_op_num_threads = 1
+            sess_options.inter_op_num_threads = 1
+            sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+            sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+            self.session = ort.InferenceSession(self.model_path, sess_options, providers=ONNX_BACKENDS)
+
+            if channel_first:
+                _, _, self.input_height, self.input_width = self.session.get_inputs()[0].shape
+            else:
+                _, self.input_height, self.input_width, _ = self.session.get_inputs()[0].shape
+
+        elif self.backend == InferModelBackend.TFLITE:
+            try:
+                import tflite_runtime.interpreter as tflite
+            except:
+                import tensorflow.lite as tflite
+
+
+            self.interpreter = tflite.Interpreter(model_path=self.model_path)
+            self.interpreter.allocate_tensors()
 
-        sess_options = ort.SessionOptions()
-        sess_options.intra_op_num_threads = 1
-        sess_options.inter_op_num_threads = 1
-        sess_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
-        sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
-        self.session = ort.InferenceSession(self.model_path, sess_options, providers=ONNX_BACKENDS)
+            # Get input and output tensors.
+            self.input_details = self.interpreter.get_input_details()
+            self.output_details = self.interpreter.get_output_details()
 
-        if channel_first:
-            _, _, self.input_height, self.input_width = self.session.get_inputs()[0].shape
+            _, self.input_height, self.input_width, _ = self.input_details[0]["shape"]
+            self.channel_first = False
         else:
-            _, self.input_height, self.input_width, _ = self.session.get_inputs()[0].shape
+            raise AttributeError(f"Backend ({self.backend}) is not supported")
 
     def preprocess(self, img):
 
         img = cv2.resize(img, (self.input_width, self.input_height))
         img_mean = np.array([127, 127, 127])
         img = (img - img_mean) / 128
 
-        img = np.transpose(img, [2, 0, 1])
+        if self.channel_first:
+            img = np.transpose(img, [2, 0, 1])
+
         img = np.expand_dims(img, axis=0)
         img = img.astype(np.float32)
 
         return img
 
-class FaceDetector(ONNXModel):
+    def _predict(self, img):
+        pass
+    
+class FaceDetector(InferModel):
+
+    def __init__(self, model_path, channel_first=True, backend=InferModelBackend.ONNX):
+        if backend != InferModelBackend.ONNX:
+            raise AttributeError(f"Backend ({backend}) is not supported for FaceDetector")
+        else:
+            super().__init__(model_path, channel_first, backend)
 
     def postprocess(self, width, height, confidences, boxes, prob_threshold, iou_threshold=0.3, top_k=-1):
         boxes = boxes[0]
@@ -82,13 +121,21 @@ def predict(self, ori_img, threshold=0.5):
 
         return boxes, probs
 
-class LandmarkAligner(ONNXModel):
-
+class LandmarkAligner(InferModel):
+
+    def _predict(self, img):
+        if self.backend == InferModelBackend.ONNX:
+            return self.session.run([], {"input": img})[0]
+        elif self.backend == InferModelBackend.TFLITE:
+            self.interpreter.set_tensor(self.input_details[0]['index'], img)
+            self.interpreter.invoke()
+            return self.interpreter.get_tensor(self.output_details[0]['index'])
+        
     def predict(self, ori_img, have_face_cls=False):
         height, width, _ = ori_img.shape
 
         img = self.preprocess(ori_img)
-        lmk = self.session.run([], {'input': img})[0]
+        lmk = self._predict(img)
 
         if have_face_cls:
             lmk, face_cls = lmk[0][:70*2].reshape((70,2)), lmk[0][70*2]
@@ -103,7 +150,7 @@ def predict(self, ori_img, have_face_cls=False):
             return lmk, face_cls
 
         return lmk
-    
+
     def predict_frame(self, frame, bbox, have_face_cls=False):
 
         fheight, fwidth, _ = frame.shape
@@ -112,7 +159,6 @@ def predict_frame(self, frame, bbox, have_face_cls=False):
         if bw*bh == 0:
             return np.zeros((70,2)), 0
 
-
         face_img = frame[lmk_box[1]:lmk_box[3], lmk_box[0]:lmk_box[2], :]
         lmk = self.predict(face_img, have_face_cls=have_face_cls)
 

fdfat/tracking/facial_sort.py

@@ -5,7 +5,7 @@
 import zmq
 import zmq.decorators as zmqd
 
-from fdfat.nn.onnx import LandmarkAligner, FaceDetector
+from fdfat.nn.infer import LandmarkAligner, FaceDetector, InferModelBackend
 from fdfat.tracking.sort import SORT
 from fdfat.utils import box_utils
 from fdfat.utils import profiler
@@ -53,7 +53,11 @@ def __init__(self, args):
         self.frame_id = 0
         self.current_faces = [] # to visualize
 
-        self.landmark = LandmarkAligner(self.args.track_landmark)
+        if self.args.track_landmark.endswith(".onnx"):
+            self.infer_backend = InferModelBackend.ONNX 
+        elif self.args.track_landmark.endswith(".tflite"):
+            self.infer_backend = InferModelBackend.TFLITE 
+        self.landmark = LandmarkAligner(self.args.track_landmark, backend=self.infer_backend)
 
     def run(self):
         self._run()