| import urllib.request
|
| import os
|
| from pathlib import Path
|
| import logging
|
|
|
| logging.basicConfig(level=logging.INFO)
|
| logger = logging.getLogger(__name__)
|
|
|
| YOLO_CONFIG = {
|
| 'weights': {
|
| 'url': 'https://github.com/AlexeyAB/darknet/releases/download/darknet_yolo_v3_optimal/yolov4.weights',
|
| 'filename': 'yolov4.weights'
|
| },
|
| 'config': {
|
| 'url': 'https://raw.githubusercontent.com/AlexeyAB/darknet/master/cfg/yolov4.cfg',
|
| 'filename': 'yolov4.cfg'
|
| }
|
| }
|
|
|
| def download_file(url: str, filename: str):
|
| """Download a file and show progress."""
|
| try:
|
| logger.info(f"Downloading {filename}...")
|
|
|
| def report_progress(count, block_size, total_size):
|
| percent = int(count * block_size * 100 / total_size)
|
| print(f"\rProgress: {percent}%", end='')
|
|
|
| filepath = Path('app/models') / filename
|
| filepath.parent.mkdir(parents=True, exist_ok=True)
|
|
|
| urllib.request.urlretrieve(
|
| url,
|
| filepath,
|
| reporthook=report_progress
|
| )
|
| print()
|
| logger.info(f"Successfully downloaded {filename}")
|
| return True
|
| except Exception as e:
|
| logger.error(f"Error downloading {filename}: {e}")
|
| return False
|
|
|
| def setup_yolo():
|
| """Download and set up YOLO files."""
|
| success = True
|
| for file_info in YOLO_CONFIG.values():
|
| if not download_file(file_info['url'], file_info['filename']):
|
| success = False
|
|
|
| if success:
|
| logger.info("YOLO setup completed successfully!")
|
| else:
|
| logger.error("YOLO setup failed. Please check the errors above.")
|
|
|
| if __name__ == "__main__":
|
| setup_yolo() |
