File size: 1,792 Bytes
d488241 |
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 |
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() # New line after progress
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() |