Skip to content

Commit ea0ded7

Browse files
committed
fix: remove tflite_runtime requirement
1 parent be71bae commit ea0ded7

File tree

1 file changed

+3
-7
lines changed

1 file changed

+3
-7
lines changed

sscma/deploy/backend/tflite_infer.py

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,11 @@
22
import torch
33
import torchvision.transforms as transforms
44

5-
from .base_infer import BaseInfer
6-
7-
from sscma.utils import lazy_import
8-
5+
import tensorflow as tf
96

7+
from .base_infer import BaseInfer
108
class TFliteInfer(BaseInfer):
119

12-
@lazy_import("tflite_runtime", install_only=True)
1310
def __init__(self, weights="sscma.tflite", device=torch.device("cpu")):
1411
super().__init__(weights=weights, device=device)
1512
self.interpreter = None
@@ -54,9 +51,8 @@ def infer(self, input_data):
5451
return results
5552

5653
def load_weights(self):
57-
from tflite_runtime.interpreter import Interpreter
5854

59-
self.interpreter = Interpreter(model_path=self.weights) # load TFLite model
55+
self.interpreter = tf.lite.Interpreter(model_path=self.weights) # load TFLite model
6056

6157
self.interpreter.allocate_tensors() # allocate
6258
self.input_details = self.interpreter.get_input_details() # inputs

0 commit comments

Comments
 (0)