Skip to content

Commit

Permalink
Merge pull request #22 from fastmachinelearning/remove_tf_from_model_…
Browse files Browse the repository at this point in the history
…scan

Update scan.py
  • Loading branch information
jmitrevs authored Dec 16, 2024
2 parents b2049bf + 760d234 commit c61e801
Showing 1 changed file with 2 additions and 6 deletions.
8 changes: 2 additions & 6 deletions scan.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,10 +48,6 @@ def scan_subdir(model_type):
if filename.endswith('.pt'):
available_models.append(filename)
data_config_avai[filename] = scan_data_config(filename[:-3])
elif model_type == 'tensorflow':
if filename.endswith('.pb'):
available_models.append(filename)
data_config_avai[filename] = scan_data_config(filename[:-3])
elif model_type == 'onnx':
if filename.endswith('.onnx'):
available_models.append(filename)
Expand All @@ -62,7 +58,7 @@ def scan_subdir(model_type):
def scan():
"""Scan all directories that contain example models and write to a file."""

available_types = ['keras', 'pytorch', 'onnx', 'tensorflow']
available_types = ['keras', 'pytorch', 'onnx']

#Dictionary to save available model list
model_dict = {}
Expand All @@ -80,4 +76,4 @@ def scan():
json.dump(model_dict, open("available_models.json", 'w' ))
json.dump(data_config_avai, open("available_data_config.json", 'w' ))

scan()
scan()

0 comments on commit c61e801

Please sign in to comment.