File size: 4,568 Bytes
0a84654 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 | """
Model Downloader - Downloads AI models from Hugging Face Hub
Automatically caches models locally after first download
FULLY PORTABLE - Works on any device with any project path
"""
from huggingface_hub import hf_hub_download
from pathlib import Path
import os
import sys
import shutil
# Detect PROJECT_ROOT dynamically
def get_project_root():
"""
Find project root by looking for config/ directory
Works regardless of where app.py is located
"""
current_path = Path(__file__).resolve() # Full path to this file
# Go up from src/utils/model_downloader.py to project root
for parent in current_path.parents:
if (parent / 'config').exists() and (parent / 'webapp').exists():
return parent
# Fallback: assume parent of src/
return current_path.parent.parent.parent
PROJECT_ROOT = get_project_root()
REPO_ID = "itsluckysharma01/NETRA-Models"
CACHE_DIR = PROJECT_ROOT / 'ai_models' # Models cached in project root
print(f"\n🔍 [Model Downloader] PROJECT_ROOT detected: {PROJECT_ROOT}")
print(f"🔍 [Model Downloader] CACHE_DIR: {CACHE_DIR}\n")
def download_model(filename):
"""
Download model from Hugging Face Hub with automatic path handling
Args:
filename: Model file path (e.g., 'ai_models/activity_recognition/violence_model.h5')
Returns:
str: Path to downloaded/cached model (absolute path)
"""
try:
# Ensure cache directory exists
CACHE_DIR.mkdir(parents=True, exist_ok=True)
# Check if model already exists in flat structure
local_path = CACHE_DIR / filename
if local_path.exists():
print(f"✅ Model cached: {filename}")
return str(local_path)
# Download from Hugging Face Hub (goes to HF cache)
print(f"📥 Downloading: {filename}")
downloaded_path = hf_hub_download(
repo_id=REPO_ID,
filename=filename,
cache_dir=str(CACHE_DIR),
local_files_only=False
)
# Copy from HF cache structure to flat ai_models/ structure
src_path = Path(downloaded_path)
# Create destination directory
local_path.parent.mkdir(parents=True, exist_ok=True)
# Copy file to flat structure
shutil.copy2(src_path, local_path)
print(f"✅ Downloaded and cached: {filename}")
return str(local_path)
except Exception as e:
print(f"❌ Error downloading {filename}: {e}")
return None
def ensure_model_exists(filename):
"""
Ensure a model exists locally, download if necessary
Args:
filename: Model file path
Returns:
bool: True if model exists or was downloaded successfully
"""
local_path = CACHE_DIR / filename
# Already exists
if local_path.exists():
return True
# Try to download
result = download_model(filename)
return result is not None
def setup_all_models():
"""Download all required models on startup"""
models = [
"ai_models/activity_recognition/violence_model.h5",
"ai_models/object_detection/yolov8n.pt",
"ai_models/pose_detection/yolo11n-pose.pt",
"ai_models/weapon_detection/best.pt",
"ai_models/analysis_models/binarycnn200.h5",
"ai_models/analysis_models/CNN93.h5",
"ai_models/analysis_models/CustomCNN.h5",
"ai_models/analysis_models/fight_detection_model.h5",
]
print("\n" + "=" * 60)
print("📥 SETTING UP AI MODELS FROM HUGGING FACE HUB")
print("=" * 60)
print(f"🔍 PROJECT_ROOT: {PROJECT_ROOT}")
print(f"🔍 CACHE_DIR: {CACHE_DIR}")
print(f"🔍 Cache exists: {CACHE_DIR.exists()}")
print("=" * 60)
downloaded = 0
cached = 0
failed = 0
for model in models:
local_path = CACHE_DIR / model
if local_path.exists():
print(f"✅ Cached: {model}")
cached += 1
else:
try:
result = download_model(model)
if result:
downloaded += 1
else:
failed += 1
except Exception as e:
print(f"⚠️ Warning: Could not load {model}")
failed += 1
print("\n" + "=" * 60)
print(f"✅ Setup Complete: {downloaded} downloaded, {cached} cached, {failed} warnings")
print(f"📁 Models should be at: {CACHE_DIR}")
print("=" * 60 + "\n") |