ds_task_recycling_project/backend/config.py

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]