Reformatting + update API metadata.

Author: UPstartDeveloper

app/main.py

@@ -1,18 +1,27 @@
 from fastapi import FastAPI, File, UploadFile
 from fire_classifier.predictor import ImagePredictor
 
-# init API
+# A: define description for the API endpoints (shown on UI)
+endpoint_metadata = [
+    {
+        "name": "classify-image",
+        "description": "Predicts the possibility that a RBG image contains fire.",
+    },
+]
+# B: init API
 app = FastAPI(
-    title="Fire Detection API",
-    description="Informs the probability that an image contains fire.",
+    title="DeepFire",
+    description="A REST API for detecting the presence of fire in an image.",
+    version="0.0.2",
+    openapi_tags=endpoint_metadata,
 )
 
-# init ML inference object
+# C: init ML inference object
 predictor_config_path = "./app/config.yaml"
 predictor = ImagePredictor.init_from_config_path(predictor_config_path)
 
 
-@app.post("/classify-image/")
+@app.post("/classify-image/", tags=["classify-image"])
 def create_upload_file(file: UploadFile = File(...)):
-    '''Predicts the possibility that a RBG image contains fire.'''
+    """Predicts the possibility that a RBG image contains fire."""
     return predictor.predict_from_file(file.file)

fire_classifier/predictor.py

@@ -19,7 +19,7 @@ def __init__(self, config: Dict[str, int or str]):
 
     @classmethod
     def init_from_config_path(cls, config_path):
-        '''Parses the config file, and instantiates a new ImagePredictor'''
+        """Parses the config file, and instantiates a new ImagePredictor"""
         # load details for setting up the model
         with open(config_path, "r") as f:
             config = yaml.load(f, yaml.SafeLoader)
@@ -28,13 +28,13 @@ def init_from_config_path(cls, config_path):
         return predictor
 
     def predict_from_array(self, arr) -> Dict[str, float]:
-        '''Returns a prediction value the sample belongs to each class.'''
+        """Returns a prediction value the sample belongs to each class."""
         pred = self.model.predict(arr[np.newaxis, ...]).ravel().tolist()
         pred = [round(x, 3) for x in pred]  # values between 0-1
         return {class_label: prob for class_label, prob in zip(self.targets, pred)}
 
     def predict_from_file(self, file_object):
-        '''Converts uploaded image to a NumPy array and classifies it.'''
+        """Converts uploaded image to a NumPy array and classifies it."""
         arr = read_from_file(file_object)
         return self.predict_from_array(arr)
 

fire_classifier/utils.py

@@ -51,12 +51,14 @@ def download_model(url, file_paths, file_sha256=None):
         else:  # need to download the model
             model_file_url = f"{url}/{model_file_path}"
             keras.utils.get_file(
-                origin=model_file_url, fname=model_file_path, 
-                cache_dir=".", cache_subdir="./model"
+                origin=model_file_url,
+                fname=model_file_path,
+                cache_dir=".",
+                cache_subdir="./model",
             )
 
 
-def load_model(url, file_paths, file_sha256=None, format='composite'):
+def load_model(url, file_paths, file_sha256=None, format="composite"):
     """
     Model reconstruction.
 
@@ -70,7 +72,7 @@ def load_model(url, file_paths, file_sha256=None, format='composite'):
         file_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.
-        format(str): currently this only supports 'composite' 
+        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.
 
@@ -79,7 +81,7 @@ def load_model(url, file_paths, file_sha256=None, format='composite'):
     """
 
     def _model_from_composite_format():
-        '''Specific to using H5 + JSON as the save format'''
+        """Specific to using H5 + JSON as the save format"""
         params_file, layers_file = file_paths
         # load the model in memory
         with open(f"./model/{layers_file}") as f:
@@ -88,14 +90,14 @@ def _model_from_composite_format():
         return model
 
     def _model_from_h5():
-        '''Specific to using a single Hadoop(H5) file'''
+        """Specific to using a single Hadoop(H5) file"""
         params_file = file_paths[0]
         return keras.models.load_model(params_file)
-    
+
     # First download the model, if needed
     download_model(url, file_paths, file_sha256)
     # load the model in memory
-    if format == 'composite':
+    if format == "composite":
         return _model_from_composite_format()
     else:  # assuming a single H5
         return _model_from_h5()