Skip to content

Commit

Permalink
(core) imporved Calculator initialization printouts by conditional …
Browse files Browse the repository at this point in the history
…coloring of the output
  • Loading branch information
amkrajewski committed Feb 24, 2024
1 parent 5d5f6a4 commit 352b3f8
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 11 deletions.
3 changes: 2 additions & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,8 @@ dependencies = [
"pymongo>=4.4",
"pysmartdl2>=2.0.0",
"dnspython",
"ruamel.yaml"
"ruamel.yaml",
"colorama"
]

[project.optional-dependencies]
Expand Down
26 changes: 16 additions & 10 deletions pysipfenn/core/pysipfenn.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from tqdm.contrib.concurrent import process_map
import natsort
from pysmartdl2 import SmartDL
from colorama import Fore, Style

# Scientific Computing Imports
import numpy as np
Expand Down Expand Up @@ -71,28 +72,28 @@ def __init__(self,
verbose: bool = True):
"""Initializes the pySIPFENN Calculator object."""
if verbose:
print('********* Initializing pySIPFENN Calculator **********')
self.verbose = verbose
print('\n********* Initializing pySIPFENN Calculator **********')
self.verbose = verbose
# dictionary with all model information
with resources.files('pysipfenn.modelsSIPFENN').joinpath('models.json').open('r') as f:
if verbose:
print(f'Loading model definitions from: {f.name}')
print(f'Loading model definitions from: {Fore.BLUE}{f.name}{Style.RESET_ALL}')
self.models = json.load(f)
# networks list
self.network_list = list(self.models.keys())
if verbose:
print(f'Found {len(self.network_list)} network definitions in models.json')
print(f'Found {Fore.BLUE}{len(self.network_list)} network definitions in models.json{Style.RESET_ALL}')
# network names
self.network_list_names = [self.models[net]['name'] for net in self.network_list]
self.network_list_available = []
self.updateModelAvailability()

self.loadedModels = {}
if autoLoad:
print(f'Loading all available models (autoLoad=True)')
print(f'Loading all available models ({Fore.BLUE}autoLoad=True{Style.RESET_ALL})')
self.loadModels()
else:
print(f'Skipping model loading (autoLoad=False)')
print(f'Skipping model loading ({Fore.BLUE}autoLoad=False{Style.RESET_ALL})')

self.prototypeLibrary = {}
self.parsePrototypeLibrary(verbose=verbose)
Expand All @@ -105,7 +106,7 @@ def __init__(self,
}
self.inputFiles = []
if verbose:
print(f'********* pySIPFENN Successfully Initialized **********')
print(f'{Fore.GREEN}********** Successfully Initialized **********{Style.RESET_ALL}')

def __str__(self):
"""Prints the status of the ``Calculator`` object."""
Expand Down Expand Up @@ -175,7 +176,12 @@ class automatically initializes with ``verbose=True``.
}
})
if verbose:
print(f'{len(self.prototypeLibrary)} prototype structures present into the prototype library.')
protoLen = len(self.prototypeLibrary)
if protoLen == 0:
print(f"{Style.DIM}No prototypes were loaded into the prototype library.{Style.RESET_ALL}")
else:
print(f"Loaded {Fore.GREEN}{protoLen} prototypes {Style.RESET_ALL}into the library.")


def appendPrototypeLibrary(self, customPath: str) -> None:
"""Parses a custom prototype library YAML file and permanently appends it into the internal prototypeLibrary
Expand Down Expand Up @@ -205,9 +211,9 @@ def updateModelAvailability(self) -> None:
for net, netName in zip(self.network_list, self.network_list_names):
if all_files.__contains__(net + '.onnx'):
detectedNets.append(net)
print('✔ ' + netName)
print(f"{Fore.GREEN}{netName}{Style.RESET_ALL}")
else:
print('x ' + netName)
print(f"{Style.DIM}x {netName}{Style.RESET_ALL}")
self.network_list_available = detectedNets

def downloadModels(self, network: str = 'all') -> None:
Expand Down

0 comments on commit 352b3f8

Please sign in to comment.