config/model_config.py

"""
Model Configuration
Centralizes all model paths and loading configurations
"""

from pathlib import Path

# Project root
PROJECT_ROOT = Path(__file__).parent.parent

# Model directory (consolidated to ai_models/)
MODELS_DIR = PROJECT_ROOT / "ai_models"

# Model paths - Organized by detection type
MODEL_PATHS = {
    'violence': {
        'model': MODELS_DIR / "violence_detection" / "violence_model.h5",
        'labels': MODELS_DIR / "violence_detection" / "lb.pickle",
        'description': 'Violence Detection Model'
    },
    'yolo': {
        'candidates': [
            MODELS_DIR / "object_detection" / "yolov8n.pt",
            MODELS_DIR / "object_detection" / "yolov8s.pt",
        ],
        'description': 'YOLO Object Detection Model'
    },
    'weapon': {
        'gun_model_candidates': [
            MODELS_DIR / "weapon_detection" / "GunDetector.pt",
            MODELS_DIR / "weapon_detection" / "best.pt",
        ],
        'person_model_candidates': [
            MODELS_DIR / "object_detection" / "yolov8n.pt",
            MODELS_DIR / "object_detection" / "yolov8s.pt",
        ],
        'description': 'Weapon & Person Detection Model'
    },
    'pose': {
        'candidates': [
            MODELS_DIR / "pose_detection" / "yolo11n-pose.pt",
        ],
        'description': 'Pose Detection Model'
    },
    'anomaly': {
        'candidates': [
            MODELS_DIR / "anomaly_detection" / "best.bin",
        ],
        'description': 'Anomaly Detection Model'
    },
    'analysis': {
        'candidates': [
            MODELS_DIR / "analysis_models" / "fight_detection_model.h5",
            MODELS_DIR / "analysis_models" / "CustomCNN150.h5",
            MODELS_DIR / "analysis_models" / "CustomCNN100.h5",
            MODELS_DIR / "analysis_models" / "CustomCNN.h5",
        ],
        'description': 'Advanced Analysis Model (Fight/Behavior Detection)'
    }
}

# Detection thresholds
DETECTION_THRESHOLDS = {
    'yolo': 0.25,
    'violence': 0.6,
    'weapon': 0.5,
    'anomaly': 0.5,
}

# Processing parameters
PROCESSING_PARAMS = {
    'frame_skip': 10,  # Process every 10th frame in video analysis
    'alert_threshold': 3,  # Consecutive frames before alert
    'max_upload_size': 500 * 1024 * 1024,  # 500MB
}


def get_model_path(model_name: str, model_type: str = None):
    """
    Get model path by name
    
    Args:
        model_name: 'violence', 'yolo', 'weapon', 'pose', 'anomaly'
        model_type: specific model variant (if applicable) - for weapon, use 'gun' or 'person'
    
    Returns:
        Path to model file or None if not found
    """
    if model_name not in MODEL_PATHS:
        return None
    
    config = MODEL_PATHS[model_name]
    
    # Handle single model file case
    if 'model' in config:
        return config['model'] if config['model'].exists() else None
    
    # Handle weapon model special case with gun and person models
    if model_name == 'weapon':
        if model_type == 'gun' and 'gun_model_candidates' in config:
            for candidate in config['gun_model_candidates']:
                if candidate.exists():
                    return candidate
        elif model_type == 'person' and 'person_model_candidates' in config:
            for candidate in config['person_model_candidates']:
                if candidate.exists():
                    return candidate
        elif not model_type:
            # Return gun model by default for weapon
            for candidate in config.get('gun_model_candidates', []):
                if candidate.exists():
                    return candidate
        return None
    
    # Handle standard candidates case
    if 'candidates' in config:
        for candidate in config['candidates']:
            if candidate.exists():
                return candidate
    
    return None


def get_all_available_models():
    """
    Get list of all available models
    
    Returns:
        dict: {model_name: is_available}
    """
    available = {}
    for model_name in MODEL_PATHS.keys():
        path = get_model_path(model_name)
        available[model_name] = path is not None
    
    return available