Ayomide
62 lines
2.0 KiB
Python
62 lines
2.0 KiB
Python
import os
|
|
from pathlib import Path
|
|
|
|
|
|
class Config:
|
|
# Server settings
|
|
HOST = os.environ.get('HOST', '0.0.0.0')
|
|
PORT = int(os.environ.get('PORT', 5000))
|
|
DEBUG = os.environ.get('FLASK_ENV') == 'development'
|
|
|
|
# File settings
|
|
UPLOAD_FOLDER = 'static/uploads'
|
|
RESULT_FOLDER = 'static/results'
|
|
MAX_FILE_SIZE = 16 * 1024 * 1024 # 16MB
|
|
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp'}
|
|
|
|
# Model settings - pre-trained model
|
|
MODEL_PATH = os.environ.get('MODEL_PATH', 'yolov8n.pt') # Auto-downloads if not exists
|
|
CONFIDENCE_THRESHOLD = float(os.environ.get('CONFIDENCE_THRESHOLD', 0.3))
|
|
IMAGE_SIZE = 416
|
|
|
|
# Hardcoded test images
|
|
TEST_IMAGES_PATH = 'training/memory'
|
|
|
|
# Logging
|
|
LOG_LEVEL = os.environ.get('LOG_LEVEL', 'INFO')
|
|
LOG_FILE = 'logs/app.log'
|
|
|
|
@classmethod
|
|
def create_directories(cls):
|
|
"""Create required directories."""
|
|
directories = [
|
|
cls.UPLOAD_FOLDER,
|
|
cls.RESULT_FOLDER,
|
|
Path(cls.LOG_FILE).parent,
|
|
]
|
|
|
|
for directory in directories:
|
|
Path(directory).mkdir(parents=True, exist_ok=True)
|
|
|
|
@classmethod
|
|
def validate_model(cls):
|
|
"""Check if model file exists or can be downloaded."""
|
|
# For pre-trained models
|
|
if cls.MODEL_PATH in ['yolov8n.pt', 'yolov8s.pt', 'yolov8m.pt', 'yolov8l.pt', 'yolov8x.pt']:
|
|
print(f"Using pre-trained model: {cls.MODEL_PATH} (will auto-download if needed)")
|
|
return
|
|
|
|
# For custom models, check if file exists
|
|
if not Path(cls.MODEL_PATH).exists():
|
|
raise FileNotFoundError(f"Model file not found: {cls.MODEL_PATH}")
|
|
|
|
@classmethod
|
|
def get_test_images(cls):
|
|
"""Get list of test images."""
|
|
test_path = Path(cls.TEST_IMAGES_PATH)
|
|
if not test_path.exists():
|
|
return []
|
|
|
|
return [f for f in test_path.iterdir()
|
|
if f.suffix.lower()[1:] in cls.ALLOWED_EXTENSIONS]
|