Spaces:
Running
Running
File size: 2,048 Bytes
6248776 |
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 |
"""
Script to download YOLOv8 model and save it to the models directory
This ensures the model is included in the repository for Hugging Face Spaces
"""
import os
from ultralytics import YOLO
def download_model(model_name="yolov8n.pt", save_dir="models"):
"""
Download a YOLOv8 model and save it to the specified directory
Args:
model_name: Name of the YOLOv8 model to download
save_dir: Directory to save the model to
"""
# Create models directory if it doesn't exist
if not os.path.exists(save_dir):
os.makedirs(save_dir)
print(f"Created directory: {save_dir}")
# Full path to save the model
model_path = os.path.join(save_dir, model_name)
# Check if model already exists
if os.path.exists(model_path):
print(f"Model already exists at: {model_path}")
# Load model to verify it's valid
try:
model = YOLO(model_path)
print("Model loaded successfully")
return model_path
except Exception as e:
print(f"Error loading existing model: {e}")
print("Downloading new model...")
# Download the model
try:
print(f"Downloading model: {model_name}")
model = YOLO(model_name)
# Save the model to the specified path
model_file = model.export(format="torchscript")
# Rename to the original model name
if os.path.exists(model_file) and model_file != model_path:
# Copy the model file to the specified path
import shutil
shutil.copy(model_file, model_path)
print(f"Model saved to: {model_path}")
return model_path
except Exception as e:
print(f"Error downloading model: {e}")
return None
if __name__ == "__main__":
# Download YOLOv8n model
model_path = download_model()
if model_path:
print(f"Model successfully downloaded to: {model_path}")
else:
print("Failed to download model") |