File size: 4,294 Bytes
80cfb2b | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 | """
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
|