fix: add model paths and class names for ONNX models in model_loader
Browse files- utils/model_loader.py +25 -0
utils/model_loader.py
CHANGED
|
@@ -15,6 +15,31 @@ import torch
|
|
| 15 |
import numpy as np
|
| 16 |
from PIL import Image
|
| 17 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 18 |
# Cache for ONNX sessions and preprocessors
|
| 19 |
_onnx_model_cache = {}
|
| 20 |
|
|
|
|
| 15 |
import numpy as np
|
| 16 |
from PIL import Image
|
| 17 |
|
| 18 |
+
|
| 19 |
+
# Model paths and class names (copied from app_mcp.py)
|
| 20 |
+
MODEL_PATHS = {
|
| 21 |
+
"model_1": "LPX55/detection-model-1-ONNX",
|
| 22 |
+
"model_2": "LPX55/detection-model-2-ONNX",
|
| 23 |
+
"model_3": "LPX55/detection-model-3-ONNX",
|
| 24 |
+
"model_4": "cmckinle/sdxl-flux-detector_v1.1",
|
| 25 |
+
"model_5": "LPX55/detection-model-5-ONNX",
|
| 26 |
+
"model_6": "LPX55/detection-model-6-ONNX",
|
| 27 |
+
"model_7": "LPX55/detection-model-7-ONNX",
|
| 28 |
+
"model_8": "aiwithoutborders-xyz/CommunityForensics-DeepfakeDet-ViT"
|
| 29 |
+
}
|
| 30 |
+
|
| 31 |
+
CLASS_NAMES = {
|
| 32 |
+
"model_1": ['artificial', 'real'],
|
| 33 |
+
"model_2": ['AI Image', 'Real Image'],
|
| 34 |
+
"model_3": ['artificial', 'human'],
|
| 35 |
+
"model_4": ['AI', 'Real'],
|
| 36 |
+
"model_5": ['Realism', 'Deepfake'],
|
| 37 |
+
"model_6": ['ai_gen', 'human'],
|
| 38 |
+
"model_7": ['Fake', 'Real'],
|
| 39 |
+
"model_8": ['Fake', 'Real'],
|
| 40 |
+
}
|
| 41 |
+
|
| 42 |
+
|
| 43 |
# Cache for ONNX sessions and preprocessors
|
| 44 |
_onnx_model_cache = {}
|
| 45 |
|