Add back supported_extensions (#137)

protectai · Apr 15, 2024 · ce8822d · ce8822d
1 parent 022985e
commit ce8822d
Showing 1 changed file with 14 additions and 0 deletions.
diff --git a/modelscan/settings.py b/modelscan/settings.py
@@ -15,31 +15,45 @@
     "scanners": {
         "modelscan.scanners.H5LambdaDetectScan": {
             "enabled": True,
+            "supported_extensions": [".h5"],
         },
         "modelscan.scanners.KerasLambdaDetectScan": {
             "enabled": True,
+            "supported_extensions": [".keras"],
         },
         "modelscan.scanners.SavedModelLambdaDetectScan": {
             "enabled": True,
+            "supported_extensions": [".pb"],
             "unsafe_keras_operators": {
                 "Lambda": "MEDIUM",
             },
         },
         "modelscan.scanners.SavedModelTensorflowOpScan": {
             "enabled": True,
+            "supported_extensions": [".pb"],
             "unsafe_tf_operators": {
                 "ReadFile": "HIGH",
                 "WriteFile": "HIGH",
             },
         },
         "modelscan.scanners.NumpyUnsafeOpScan": {
             "enabled": True,
+            "supported_extensions": [".npy"],
         },
         "modelscan.scanners.PickleUnsafeOpScan": {
             "enabled": True,
+            "supported_extensions": [
+                ".pkl",
+                ".pickle",
+                ".joblib",
+                ".dill",
+                ".dat",
+                ".data",
+            ],
         },
         "modelscan.scanners.PyTorchUnsafeOpScan": {
             "enabled": True,
+            "supported_extensions": [".bin", ".pt", ".pth", ".ckpt"],
         },
     },
     "middlewares": {