From 2213af2220f1d2c77ab6de67dc53c10f9c0eb25b Mon Sep 17 00:00:00 2001
From: Sam Washko <sam@protectai.com>
Date: Wed, 10 Apr 2024 17:11:27 -0700
Subject: [PATCH] tensorflow extras

---
 README.md                              | 7 ++++++-
 modelscan/__init__.py                  | 1 +
 modelscan/scanners/saved_model/scan.py | 4 +++-
 3 files changed, 10 insertions(+), 2 deletions(-)

diff --git a/README.md b/README.md
index 57c9027..01a6069 100644
--- a/README.md
+++ b/README.md
@@ -31,7 +31,7 @@ pip install modelscan
 With it installed, scan a model:
 
 ```bash
-modelscan -p /path/to/model_file.h5
+modelscan -p /path/to/model_file.pkl
 ```
 
 ## Why You Should Scan Models
@@ -105,6 +105,11 @@ or `pyproject.toml` like this:
 modelscan = ">=0.1.1"
 ```
 
+Scanners for Tensorflow or HD5 formatted models require installation with extras:
+```bash
+pip install 'modelscan[ tensorflow, h5py ]'
+```
+
 ### Using ModelScan via CLI
 
 ModelScan supports the following arguments via the CLI:
diff --git a/modelscan/__init__.py b/modelscan/__init__.py
index b420341..9868bc5 100644
--- a/modelscan/__init__.py
+++ b/modelscan/__init__.py
@@ -1,4 +1,5 @@
 """CLI for scanning models"""
+
 import logging
 
 from modelscan._version import __version__
diff --git a/modelscan/scanners/saved_model/scan.py b/modelscan/scanners/saved_model/scan.py
index 58d1414..c419bd2 100644
--- a/modelscan/scanners/saved_model/scan.py
+++ b/modelscan/scanners/saved_model/scan.py
@@ -65,7 +65,9 @@ def _check_for_unsafe_tf_keras_operator(
         unsafe_operators: Dict[str, Any],
     ) -> ScanResults:
         issues: List[Issue] = []
-        all_operators = tensorflow.raw_ops.__dict__.keys()
+        all_operators = (
+            tensorflow.raw_ops.__dict__.keys() if tensorflow_installed else []
+        )
         all_safe_operators = [
             operator for operator in list(all_operators) if operator[0] != "_"
         ]