Refactor classifer utils.

Author: UPstartDeveloper

.gitignore

@@ -11,7 +11,7 @@ __pycache__/
 .vscode
 
 # Keras Models
-model/*
+*model/*
 
 # Distribution / packaging
 .Python

app/main.py

@@ -1,17 +1,17 @@
 from fastapi import FastAPI, File, UploadFile
 from fire_classifier.predictor import ImagePredictor
-from fire_classifier.util.api import API
+from app.settings import API_SETTINGS
 
 # A: init API
 app = FastAPI(
-    title=API["title"],
-    description=API["description"],
-    version=API["version"],
-    openapi_tags=API["endpoints"],
+    title=API_SETTINGS["title"],
+    description=API_SETTINGS["description"],
+    version=API_SETTINGS["version"],
+    openapi_tags=API_SETTINGS["openapi_tags"],
 )
 
 # B: init ML inference object, and the routes
-predictor_config_path = API["config_path"]
+predictor_config_path = API_SETTINGS["predictor_config_path"]
 predictor = ImagePredictor.init_from_config_path(predictor_config_path)
 
 

fire_classifier/util/api.py app/settings.pyfire_classifier/util/api.py renamed to app/settings.py

@@ -10,22 +10,22 @@
 # These are intended shown on the UI, related only to a specific endpoint.
 # Note: the value in the "name" field should match what goes in the
 #       "tags" parameter of the corresponding app route in main.py!!
-API_ENDPOINTS = (
+API_ENDPOINT_DATA = (
     {
         "name": "Detect Fire",
         "description": "Predicts the possibility that a color image contains fire.",
     },
 )
 
 # Tells the app how to find the config.yaml (for running ML inference)
-BASE_DIR = Path(__file__).resolve().parent.parent.parent
+BASE_DIR = Path(__file__).resolve().parent.parent
 CONFIG_PATH = os.path.join(BASE_DIR, "app", "config.yaml")
 
 # Wraps all the API metadata as one dictionary
-API = {
+API_SETTINGS = {
     "title": API_TITLE,
     "description": API_DESCRIPTION,
     "version": API_VERSION,
-    "endpoints": API_ENDPOINTS,
-    "config_path": CONFIG_PATH,
+    "openapi_tags": API_ENDPOINT_DATA,
+    "predictor_config_path": CONFIG_PATH,
 }

fire_classifier/predictor.py

@@ -4,17 +4,15 @@
 from typing import Dict
 import yaml
 
-from fire_classifier.preprocessing_utilities import read_from_file
-from fire_classifier.utils import load_model
+from fire_classifier.util import preprocessing
+from fire_classifier.util.model import ModelUtility
 
 
 class ImagePredictor:
     def __init__(self, config: Dict[str, int or str]):
         self.model_paths = config["model_file_paths"]
         self.resize_size = config["resize_shape"]
-        self.model = load_model(
-            config["base_model_url"], self.model_paths, config["model_sha256"]
-        )
+        self.model = ModelUtility.reconstruct_model(config)
         self.targets = config["targets"]
 
     @classmethod
@@ -35,7 +33,7 @@ def predict_from_array(self, arr) -> Dict[str, float]:
 
     def predict_from_file(self, file_object):
         """Converts uploaded image to a NumPy array and classifies it."""
-        arr = read_from_file(file_object)
+        arr = preprocessing.read_from_file(file_object)
         return self.predict_from_array(arr)
 
 

fire_classifier/util/__init__.py

@@ -0,0 +1,119 @@
+import hashlib
+import os
+from typing import Dict
+from tensorflow import keras
+
+
+class ModelUtility:
+    def __init__(self, config: Dict[str, str]):
+        """
+        Instaniates a new object using data needed to load in the model.
+
+        Args:
+        config(dict): contains the following fields of interest:
+            base_model_url(str): the base url where the files are located
+            model_file_paths(list): collection of all the files needed to
+                                         eventually load the model
+            model_sha256(str): the supposed hash of one of the files
+                               we need to download. Checked against the
+                               one we may already have in the codebase.
+        """
+        self.url = config["base_model_url"]
+        self.file_paths = config["model_file_paths"]
+        self.file_sha256 = None
+        if config["model_sha256"] is not None:
+            self.file_sha256 = config["model_sha256"]
+
+    @classmethod
+    def reconstruct_model(cls, config):
+        '''Make a new instance, and load in the model straightaway.'''
+        model_utility = cls(config)
+        # detect save format
+        save_format = 'composite'
+        if config["model_file_paths"] and len(config["model_file_paths"]) == 1:
+            save_format = 'h5'
+        # load the model
+        return model_utility.load_model(save_format)
+
+    def get_hash(self, filename):
+        """
+        Computes the SHA256 hash of a given file.
+
+        This can then be used to ensure the model file(s) downloaded
+        in this codebase are not corrupted.
+
+        Args:
+            filename(str): the name of the file
+
+        Returns:
+            bytes-like object
+        """
+        sha256_hash = hashlib.sha256()
+        with open(filename, "rb") as f:
+            for byte_block in iter(lambda: f.read(4096), b""):
+                sha256_hash.update(byte_block)
+
+        return sha256_hash.hexdigest()
+
+    def download_model(self):
+        """
+        Downloads the model files in memory.
+
+        This will first check if the files are already present,
+        and not corrupted, before downloading from the address
+        specified in config.yaml.
+
+        Returns:
+            None
+        """
+        # Download only the model files that are needed
+        for model_file_path in self.file_paths:
+            if os.path.exists(model_file_path):
+                if self.get_hash(model_file_path) == self.file_sha256:
+                    print(f"File already exists: {model_file_path}")
+            else:  # need to download the model
+                model_file_url = f"{self.url}/{model_file_path}"
+                keras.utils.get_file(
+                    origin=model_file_url,
+                    fname=model_file_path,
+                    cache_dir=".",
+                    cache_subdir="./model",
+                )
+
+    def load_model(self, format="composite"):
+        """
+        Model reconstruction.
+
+        This will first load the model in memory using the given files
+        and save format
+
+        Args:
+            format(str): currently this only supports 'composite'
+                        (which is for when the model is saved using a H5 + JSON)
+                        or 'h5' as the save format of the model.
+
+        Returns:
+            keras.Model object
+        """
+
+        def _model_from_composite_format():
+            """Specific to using H5 + JSON as the save format"""
+            params_file, layers_file = self.file_paths
+            # load the model in memory
+            with open(f"./model/{layers_file}") as f:
+                model = keras.models.model_from_json(f.read())  # build the layers
+                model.load_weights(f"./model/{params_file}")  # load weights + biases
+            return model
+
+        def _model_from_h5():
+            """Specific to using a single Hadoop(H5) file"""
+            params_file = self.file_paths[0]
+            return keras.models.load_model(params_file)
+
+        # First download the model, if needed
+        self.download_model()
+        # load the model in memory
+        if format == "composite":
+            return _model_from_composite_format()
+        else:  # assuming a single H5
+            return _model_from_h5()