LPX55 commited on
Commit
a4381af
·
1 Parent(s): 0adef02

fix: add model paths and class names for ONNX models in model_loader

Browse files
Files changed (1) hide show
  1. 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