separate models and add image_captioning

.gitignore

@@ -0,0 +1,5 @@
+__pycache__
+.vscode
+
+*.jpg
+*.png

app.py

@@ -1,74 +1,38 @@
 from pathlib import Path
 
 import numpy as np
-import tensorflow as tf
 import gradio as gr
 
-import cv2
-import keras
-from keras import Sequential
-from keras.applications.resnet50 import ResNet50
-from keras.layers import Flatten, Dense
+from lib.image_captioning import ImageCaptioning
+from lib.pace_model import PaceModel
 
 pace_model_weights_path = (Path.cwd() / "models" / "pace_model_weights.h5").resolve()
 resnet50_tf_model_weights_path = (Path.cwd() / "models" / "resnet50_weights_tf_dim_ordering_tf_kernels_notop.h5")
 height, width, channels = (224, 224, 3)
 
-class PaceModel:
-    def __init__(self, height, width, channels):
-        self.resnet_model = Sequential()
-        self.height = height
-        self.width = width
-        self.channels = channels
-        self.class_names = ["Fast", "Medium", "Slow"]
-
-        self.create_base_model()
-        self.create_architecture()
-    
-    def create_base_model(self):
-        self.base_model = ResNet50(
-            include_top=False,
-            input_shape=(self.height, self.width, self.channels),
-            pooling="avg",
-            classes=211,
-            weights="imagenet"
-        )
-        self.base_model.load_weights(resnet50_tf_model_weights_path)
-
-        for layer in self.base_model.layers:
-            layer.trainable = False
+class AudioPalette:
+    def __init__(self):
+        self.pace_model = PaceModel(height, width, channels, resnet50_tf_model_weights_path, pace_model_weights_path)
+        self.image_captioning = ImageCaptioning()
     
-    def create_architecture(self):
-        self.resnet_model.add(self.base_model)
-        self.resnet_model.add(Flatten())
-        self.resnet_model.add(Dense(1024, activation="relu"))
-        self.resnet_model.add(Dense(256, activation="relu"))
-        self.resnet_model.add(Dense(3, activation="softmax"))
-
-        self.resnet_model.load_weights(pace_model_weights_path)
-            
-    def predict(self, input_image: np.ndarray):
-        resized_image = cv2.resize(input_image, (self.height, self.width))
-        image = np.expand_dims(resized_image, axis=0)
-
-        prediction = self.resnet_model.predict(image)
-        print(prediction, np.argmax(prediction))
-        return self.class_names[np.argmax(prediction)]
+    def generate(self, input_image_path):
+        generated_text = self.image_captioning.query(input_image_path)[0].get("generated_text")
+        return self.pace_model.predict(input_image_path) + " - " + generated_text
 
 def main():
-    model = PaceModel(height, width, channels)
+    model = AudioPalette()
 
     demo = gr.Interface(
-        fn=model.predict,
+        fn=model.generate,
         inputs=gr.Image(
-            type="numpy",
+            type="filepath",
             label="Upload an image",
             show_label=True,
             container=True
         ),
         outputs=gr.Textbox(
             lines=1,
-            placeholder="Fast | Medium | Slow",
+            placeholder="Pace of the image and the caption",
             label="Pace of the image",
             show_label=True,
             container=True,

lib/image_captioning.py

@@ -0,0 +1,27 @@
+import os
+
+import cv2
+import requests
+
+class ImageCaptioning:
+    """
+    Performing an API call to BLIP's huggingface inference API
+    """
+    def __init__(self):
+        self.api_endpoint = os.environ["blip_api_url"]
+        self.org_token = os.environ["auth_token"]
+        self.headers = { "Authorization": f"Bearer {self.org_token}" }
+    
+    def read_image(self, image_path):
+        with open(image_path, "rb") as f:
+            data = f.read()
+        
+        return data
+
+    def query(self, image_path: str):
+        response = requests.post(
+            self.api_endpoint,
+            headers=self.headers,
+            data=self.read_image(image_path)
+        )
+        return response.json()

lib/pace_model.py

@@ -0,0 +1,55 @@
+import numpy as np
+import tensorflow as tf
+
+import cv2
+import keras
+from keras import Sequential
+from keras.applications.resnet50 import ResNet50
+from keras.layers import Flatten, Dense
+
+class PaceModel:
+    """
+    The pace model which uses ResNet50's architecture as base and builds upon by adding further layers to determine the pace of an image.
+    """
+    def __init__(self, height, width, channels, resnet50_tf_model_weights_path, pace_model_weights_path):
+        self.resnet_model = Sequential()
+        self.height = height
+        self.width = width
+        self.channels = channels
+        self.class_names = ["Fast", "Medium", "Slow"]
+        self.resnet50_tf_model_weights_path = resnet50_tf_model_weights_path
+        self.pace_model_weights_path = pace_model_weights_path
+
+        self.create_base_model()
+        self.create_architecture()
+    
+    def create_base_model(self):
+        self.base_model = ResNet50(
+            include_top=False,
+            input_shape=(self.height, self.width, self.channels),
+            pooling="avg",
+            classes=211,
+            weights="imagenet"
+        )
+        self.base_model.load_weights(self.resnet50_tf_model_weights_path)
+
+        for layer in self.base_model.layers:
+            layer.trainable = False
+    
+    def create_architecture(self):
+        self.resnet_model.add(self.base_model)
+        self.resnet_model.add(Flatten())
+        self.resnet_model.add(Dense(1024, activation="relu"))
+        self.resnet_model.add(Dense(256, activation="relu"))
+        self.resnet_model.add(Dense(3, activation="softmax"))
+
+        self.resnet_model.load_weights(self.pace_model_weights_path)
+            
+    def predict(self, input_image_path: str):
+        input_image = cv2.imread(input_image_path)
+        resized_image = cv2.resize(input_image, (self.height, self.width))
+        image = np.expand_dims(resized_image, axis=0)
+
+        prediction = self.resnet_model.predict(image)
+        print(prediction, np.argmax(prediction))
+        return self.class_names[np.argmax(prediction)]