Oill_split / backend /utils /model_fetcher.py
Utkarshres32's picture
Initial commit for Hugging Face deployment
998bc6e
import os
import urllib.request
import sys
def download_if_missing():
"""
Checks if the model weights exist locally.
If not, attempts to download them from the URL provided in environment variables.
"""
model_url = os.getenv("MODEL_URL")
# Default local path relative to backend directory
model_path = os.path.join(os.path.dirname(__file__), "../../model/saved_models/oil_spill_unet_best.keras")
# Ensure directory exists
os.makedirs(os.path.dirname(model_path), exist_ok=True)
if not os.path.exists(model_path):
if not model_url or model_url == "https://your-r2-public-url/oil_spill_unet_best.keras":
print("⚠️ MODEL_URL not configured. Skipping download.")
print("πŸ’‘ Please provide a valid URL in your docker-compose.yml to enable automated fetching.")
return
print(f"πŸš€ Model weights missing. Initializing download from: {model_url}")
try:
# Using urllib as it's built-in and sufficient for this simple GET
urllib.request.urlretrieve(model_url, model_path)
print("βœ… Model weights downloaded successfully.")
except Exception as e:
print(f"❌ Failed to download model: {e}")
print("πŸ’‘ Fallback: The system will use an untrained stub model for now.")
else:
print("πŸ“¦ Production model weights detected locally. Skipping download.")
if __name__ == "__main__":
# Test run
download_if_missing()