Spaces:
Running
Running
| """ | |
| 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") |