Replace tflite with ai-edge-litert

Author: omad

docker/requirements.txt

@@ -1,5 +1,8 @@
 
 --extra-index-url https://packages.dea.ga.gov.au/
+
+# For ML
+ai-edge-litert
 datacube-ows>=1.9
 datacube[performance,s3]>=1.9.5
 eodatasets3>1.9
@@ -16,7 +19,4 @@ odc-stac>=0.4.0
 
 # odc-stac is in PyPI
 odc-stats[ows]
-
-# For ML
-tflite-runtime
 tl2cgen

odc/stats/plugins/_worker.py

@@ -5,7 +5,7 @@
 """
 
 from dask.distributed import WorkerPlugin, get_worker
-import tflite_runtime.interpreter as tflite
+from ai_edge_litert.interpreter import Interpreter
 import tl2cgen
 import threading
 import logging
@@ -25,7 +25,7 @@ def get_interpreter(self):
         worker = get_worker()
         thread_id = threading.get_ident()
         if thread_id not in worker.interpreters:
-            interpreter = tflite.Interpreter(model_path=self.model_path)
+            interpreter = Interpreter(model_path=self.model_path)
             interpreter.allocate_tensors()
             worker.interpreters[thread_id] = interpreter
             self._log.info(

setup.cfg

@@ -42,7 +42,7 @@ install_requires =
     fiona
     rasterio>=1.3.2
     s3fs
-    tflite-runtime
+    ai-edge-litert
     tl2cgen
 
 [options.entry_points]

tests/requirements.txt

@@ -1,4 +1,5 @@
 --extra-index-url https://packages.dea.ga.gov.au/
+ai-edge-litert
 datacube>=1.9.5
 datacube-ows>=1.9
 # for pytest-depends
@@ -22,7 +23,6 @@ pytest-cov
 pytest-httpserver
 pytest-timeout
 s3fs
-tflite-runtime
 tl2cgen
 
 # patch image