|
|
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() |
