fix: update EfficientNet import issue

Kohulan · Jul 31, 2023 · 758a0f4 · 758a0f4
1 parent 664646d
commit 758a0f4
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 3 deletions.
diff --git a/DECIMER/Efficient_Net_encoder.py b/DECIMER/Efficient_Net_encoder.py
@@ -1,6 +1,7 @@
 # EfficientNet-V2 config
 import tensorflow as tf
 import DECIMER.efficientnetv2 as efficientnetv2
+from DECIMER.efficientnetv2 import effnetv2_model
 
 BATCH_SIZE_DEBUG = 2
 MODEL = "efficientnetv2-m"  # @param
@@ -16,7 +17,7 @@ def get_efficientnetv2_backbone(
             "\n...At this time we only want to use the raw "
             "(no pretraining), headless, features with no pooling ...\n"
         )
-    backbone = efficientnetv2.effnetv2_model.EffNetV2Model(model_name=model_name)
+    backbone = effnetv2_model.EffNetV2Model(model_name=model_name)
     backbone(
         tf.ones((BATCH_SIZE_DEBUG, *input_shape)), training=False, features_only=True
     )

diff --git a/DECIMER/Predictor_usingCheckpoints.py b/DECIMER/Predictor_usingCheckpoints.py
@@ -13,11 +13,12 @@
 print(tf.__version__)
 
 # Set GPU
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+"""
 gpus = tf.config.experimental.list_physical_devices("GPU")
 for gpu in gpus:
     tf.config.experimental.set_memory_growth(gpu, True)
-
+"""
 # Set path
 default_path = pystow.join("DECIMER-V2")
 

diff --git a/setup.py b/setup.py
@@ -33,6 +33,7 @@
         "pillow-heif",
         "efficientnet",
         "selfies",
+        "pyyaml",
         "tensorflow-addons",
     ],
     package_data={"DECIMER": ["repack/*.*", "efficientnetv2/*.*", "Utils/*.*"]},