Omarelrayes commited on
Commit
6e9cd12
Β·
verified Β·
1 Parent(s): 2b9a733

Update app/configs.py

Browse files
Files changed (1) hide show
  1. app/configs.py +50 -35
app/configs.py CHANGED
@@ -3,10 +3,11 @@ import os
3
  os.environ["PYTHONPATH"] = "/app:" + os.environ.get("PYTHONPATH", "")
4
 
5
  import tensorflow as tf
6
- import keras
7
  from pathlib import Path
8
  from typing import Dict, Any, List
9
  from huggingface_hub import hf_hub_download
 
 
10
 
11
  # ======================
12
  # πŸ”₯ Import HF Plugin
@@ -30,17 +31,14 @@ os.environ["HF_TOKEN"] = HF_TOKEN
30
  # ======================
31
  HF_REPO = "omarelrayes/mlflow-artifacts"
32
 
33
- # πŸ”₯ Ψ§Ω„Ω…Ψ³Ψ§Ψ±Ψ§Ψͺ Ψ§Ω„Ψ«Ψ§Ψ¨ΨͺΨ© (Ω…Ω† Ψ§Ω„ΨͺΨ³Ψ¬ΩŠΩ„ Ψ§Ω„Ψ³Ψ§Ψ¨Ω‚)
34
- CLASSIFICATION_MODEL_PATH = "1/models/m-2fa9b990eeac4e18b925616422228223/artifacts/data/model.keras"
35
- SEGMENTATION_MODEL_PATH = "2/models/m-21ba32f804234073952d1a8e590958ec/artifacts/data/model.keras"
36
 
37
  print("="*60)
38
  print("πŸ”§ Configuration:")
39
  print(f" - HF Repository: {HF_REPO}")
40
- print(f" - Classification Path: {CLASSIFICATION_MODEL_PATH}")
41
- print(f" - Segmentation Path: {SEGMENTATION_MODEL_PATH}")
42
  print(f" - TensorFlow version: {tf.__version__}")
43
- print(f" - Keras version: {keras.__version__}")
44
  print("="*60)
45
 
46
  # ======================
@@ -49,6 +47,37 @@ print("="*60)
49
  _classification_model = None
50
  _segmentation_model = None
51
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
52
  # ======================
53
  # Classification Model Loader
54
  # ======================
@@ -61,23 +90,16 @@ def get_classification_model():
61
  print("="*60)
62
 
63
  try:
64
- # πŸ”₯ Download from HuggingFace Hub
65
- model_file = hf_hub_download(
66
- repo_id=HF_REPO,
67
- filename=CLASSIFICATION_MODEL_PATH,
68
- repo_type="model",
69
- token=HF_TOKEN,
70
- cache_dir="/tmp/hf_cache"
71
  )
72
 
73
- print(f"πŸ“₯ Model downloaded to: {model_file}")
74
 
75
- # πŸ”₯ πŸ”₯ πŸ”₯ CRITICAL: Use keras.saving.load_model with safe_mode=False
76
- _classification_model = keras.saving.load_model(
77
- model_file,
78
- compile=False,
79
- safe_mode=False # πŸ”₯ Ψͺجاوز Ω…Ψ΄ΩƒΩ„Ψ© quantization_config
80
- )
81
 
82
  print("βœ… Classification Model Loaded Successfully!")
83
  except Exception as e:
@@ -100,23 +122,16 @@ def get_segmentation_model():
100
  print("="*60)
101
 
102
  try:
103
- # πŸ”₯ Download from HuggingFace Hub
104
- model_file = hf_hub_download(
105
- repo_id=HF_REPO,
106
- filename=SEGMENTATION_MODEL_PATH,
107
- repo_type="model",
108
- token=HF_TOKEN,
109
- cache_dir="/tmp/hf_cache"
110
  )
111
 
112
- print(f"πŸ“₯ Model downloaded to: {model_file}")
113
 
114
- # πŸ”₯ πŸ”₯ πŸ”₯ CRITICAL: Use keras.saving.load_model with safe_mode=False
115
- _segmentation_model = keras.saving.load_model(
116
- model_file,
117
- compile=False,
118
- safe_mode=False # πŸ”₯ Ψͺجاوز Ω…Ψ΄ΩƒΩ„Ψ© quantization_config
119
- )
120
 
121
  print("βœ… Segmentation Model Loaded Successfully!")
122
  except Exception as e:
 
3
  os.environ["PYTHONPATH"] = "/app:" + os.environ.get("PYTHONPATH", "")
4
 
5
  import tensorflow as tf
 
6
  from pathlib import Path
7
  from typing import Dict, Any, List
8
  from huggingface_hub import hf_hub_download
9
+ import zipfile
10
+ import shutil
11
 
12
  # ======================
13
  # πŸ”₯ Import HF Plugin
 
31
  # ======================
32
  HF_REPO = "omarelrayes/mlflow-artifacts"
33
 
34
+ # πŸ”₯ Ω…Ψ³Ψ§Ψ±Ψ§Ψͺ SavedModel zip files
35
+ CLASSIFIER_ZIP_PATH = "models/classifier_savedmodel.zip"
36
+ SEGMENTER_ZIP_PATH = "models/segmenter_savedmodel.zip"
37
 
38
  print("="*60)
39
  print("πŸ”§ Configuration:")
40
  print(f" - HF Repository: {HF_REPO}")
 
 
41
  print(f" - TensorFlow version: {tf.__version__}")
 
42
  print("="*60)
43
 
44
  # ======================
 
47
  _classification_model = None
48
  _segmentation_model = None
49
 
50
+ # ======================
51
+ # Helper: Download and extract SavedModel
52
+ # ======================
53
+ def download_and_extract_savedmodel(zip_path_in_repo, extract_dir):
54
+ """Download zip from HF Hub and extract SavedModel"""
55
+ print(f"πŸ“₯ Downloading {zip_path_in_repo}...")
56
+
57
+ # Download zip
58
+ zip_file = hf_hub_download(
59
+ repo_id=HF_REPO,
60
+ filename=zip_path_in_repo,
61
+ repo_type="model",
62
+ token=HF_TOKEN,
63
+ cache_dir="/tmp/hf_cache"
64
+ )
65
+
66
+ print(f"πŸ“¦ Downloaded to: {zip_file}")
67
+
68
+ # Extract
69
+ if os.path.exists(extract_dir):
70
+ shutil.rmtree(extract_dir)
71
+
72
+ os.makedirs(extract_dir, exist_ok=True)
73
+ print(f"πŸ“‚ Extracting to: {extract_dir}")
74
+
75
+ with zipfile.ZipFile(zip_file, 'r') as zipf:
76
+ zipf.extractall(extract_dir)
77
+
78
+ print(f"βœ… Extracted successfully!")
79
+ return extract_dir
80
+
81
  # ======================
82
  # Classification Model Loader
83
  # ======================
 
90
  print("="*60)
91
 
92
  try:
93
+ # πŸ”₯ Download and extract SavedModel
94
+ extract_dir = download_and_extract_savedmodel(
95
+ CLASSIFIER_ZIP_PATH,
96
+ "/tmp/savedmodels/classifier"
 
 
 
97
  )
98
 
99
+ print(f"πŸ€– Loading SavedModel from: {extract_dir}")
100
 
101
+ # πŸ”₯ Load SavedModel
102
+ _classification_model = tf.saved_model.load(extract_dir)
 
 
 
 
103
 
104
  print("βœ… Classification Model Loaded Successfully!")
105
  except Exception as e:
 
122
  print("="*60)
123
 
124
  try:
125
+ # πŸ”₯ Download and extract SavedModel
126
+ extract_dir = download_and_extract_savedmodel(
127
+ SEGMENTER_ZIP_PATH,
128
+ "/tmp/savedmodels/segmenter"
 
 
 
129
  )
130
 
131
+ print(f"πŸ€– Loading SavedModel from: {extract_dir}")
132
 
133
+ # πŸ”₯ Load SavedModel
134
+ _segmentation_model = tf.saved_model.load(extract_dir)
 
 
 
 
135
 
136
  print("βœ… Segmentation Model Loaded Successfully!")
137
  except Exception as e: