Spaces:
Sleeping
Sleeping
| import os | |
| import gdown | |
| from dotenv import load_dotenv | |
| load_dotenv() | |
| def extract_file_id_from_url(url): | |
| """Extract file ID from Google Drive URL""" | |
| if "drive.google.com" in url: | |
| if "/file/d/" in url: | |
| return url.split("/file/d/")[1].split("/")[0] | |
| elif "id=" in url: | |
| return url.split("id=")[1].split("&")[0] | |
| return url | |
| def get_model_urls(): | |
| """Get model URLs from environment variables""" | |
| return { | |
| "../models/MLII-latest.keras": os.getenv("ECG_MODEL_URL", ""), | |
| "../models/pcg_model.h5": os.getenv("PCG_MODEL_URL", ""), | |
| "../models/emg_model.h5": os.getenv("EMG_MODEL_URL", ""), | |
| "../models/vag_feature_classifier.pkl": os.getenv("VAG_MODEL_URL", "") | |
| } | |
| def download_from_gdrive(url, output_path): | |
| """Download file from Google Drive using gdown""" | |
| os.makedirs(os.path.dirname(output_path), exist_ok=True) | |
| file_id = extract_file_id_from_url(url) | |
| download_url = f"https://drive.google.com/uc?export=download&id={file_id}" | |
| gdown.download(download_url, output_path, quiet=False) | |
| def ensure_models_downloaded(): | |
| """Download models if they don't exist locally""" | |
| model_urls = get_model_urls() | |
| for local_path, url in model_urls.items(): | |
| if not url: | |
| print(f"⚠️ No URL found for {local_path}") | |
| continue | |
| if not os.path.exists(local_path): | |
| print(f"Downloading {local_path}...") | |
| try: | |
| download_from_gdrive(url, local_path) | |
| print(f"✅ Downloaded {local_path}") | |
| except Exception as e: | |
| print(f"❌ Failed to download {local_path}: {e}") | |
| else: | |
| print(f"✅ {local_path} already exists") | |
| if __name__ == "__main__": | |
| ensure_models_downloaded() |