Support tensorflow.lite in addition to tflite_runtime.

pkejval · Oct 14, 2022 · 080c71e · 080c71e
1 parent 76cce4e
commit 080c71e
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 7 deletions.
diff --git a/uldlib/captcha.py b/uldlib/captcha.py
@@ -1,12 +1,13 @@
 from abc import abstractmethod
 import threading
 import time
+import importlib
 from typing import Dict
 import requests
 from PIL import Image
 from io import BytesIO
-from uldlib.frontend import Frontend
 
+from uldlib.frontend import Frontend
 from uldlib.utils import LogLevel
 
 
@@ -105,7 +106,16 @@ def __init__(self, model_path, model_url, frontend):
 
         from urllib.request import urlretrieve
         import os
-        import tflite_runtime.interpreter as tflite
+
+        tfull_available = importlib.util.find_spec('tensorflow.lite')
+        tflite_available = importlib.util.find_spec('tflite_runtime')
+
+        if tfull_available:
+            import tensorflow.lite as tflite 
+        elif tflite_runtime:
+            import tflite_runtime.interpreter as tflite
+        else:
+            raise ImportError('No tensorflow.lite or tflite_runtime available.')
 
         def reporthook(blocknum, block_size, total_size):
             """

diff --git a/uldlib/cmd.py b/uldlib/cmd.py
@@ -34,26 +34,30 @@ def run():
     # TODO: implement other frontends and allow to choose from them
     frontend = ConsoleFrontend()
 
+    tfull_available = importlib.util.find_spec('tensorflow.lite')
     tflite_available = importlib.util.find_spec('tflite_runtime')
     tkinter_available = importlib.util.find_spec('tkinter')
 
     # Autodetection
     if not args.auto_captcha and not args.manual_captcha:
-        if tflite_available:
+        if tfull_available:
+            frontend.main_log("[Autodetect] tensorflow.lite available, using --auto-captcha")
+            args.auto_captcha = True
+        elif tflite_available:
             frontend.main_log("[Autodetect] tflite_runtime available, using --auto-captcha")
             args.auto_captcha = True
         elif tkinter_available:
             frontend.main_log("[Autodetect] tkinter available, using --manual-captcha")
             args.manual_captcha = True
         else:
             frontend.main_log(
-                "[Autodetect] WARNING: No tflite_runtime and no tkinter available, cannot solve CAPTCHA (only direct download available)",
-                level=LogLevel.WARNING
+                "[Autodetect] WARNING: No tensorflow.lite or tflite_runtime and no tkinter available, cannot solve CAPTCHA (only direct download available)",
+                level = LogLevel.WARNING
             )
 
     if args.auto_captcha:
-        if not tflite_available:
-            frontend.main_log('ERROR: --auto-captcha used but tflite_runtime not available', level=LogLevel.ERROR)
+        if not(tfull_available or tflite_available):
+            frontend.main_log('ERROR: --auto-captcha used but neither tensorflow.lite nor tflite_runtime are available', level=LogLevel.ERROR)
             sys.exit(1)
 
         model_path = path.join(__path__[0], "model.tflite")