Spaces:
Running on Zero
Running on Zero
| """ | |
| Run this once after cloning the repo to download required models. | |
| Usage: python setup_models.py | |
| """ | |
| import os | |
| import shutil | |
| import gdown | |
| from ultralytics import YOLO | |
| os.makedirs("models", exist_ok=True) | |
| # Download yolo11n.pt | |
| print("Downloading yolo11n.pt...") | |
| YOLO("yolo11n.pt") | |
| if os.path.exists("yolo11n.pt"): | |
| shutil.move("yolo11n.pt", "models/yolo11n.pt") | |
| print("yolo11n.pt ready at models/yolo11n.pt") | |
| # Download emergency.pt from Google Drive | |
| EMERGENCY_GDRIVE_URL = "https://drive.google.com/uc?id=18K2Vfv3ADBPbxU-XOMjr3ld2lX9sP6ti" | |
| EMERGENCY_PATH = "models/emergency.pt" | |
| if not os.path.exists(EMERGENCY_PATH): | |
| print("\nDownloading emergency.pt from Google Drive...") | |
| gdown.download(EMERGENCY_GDRIVE_URL, EMERGENCY_PATH, quiet=False) | |
| print("emergency.pt ready at models/emergency.pt") | |
| else: | |
| print("emergency.pt already exists.") | |