LPX55
commited on
Commit
·
a4381af
1
Parent(s):
0adef02
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 |
|