update image input to pillow

app.py

@@ -3,6 +3,8 @@ from pathlib import Path
 import numpy as np
 import gradio as gr
 
+import PIL
+
 from lib.image_captioning import ImageCaptioning
 from lib.pace_model import PaceModel
 
@@ -15,18 +17,18 @@ class AudioPalette:
         self.pace_model = PaceModel(height, width, channels, resnet50_tf_model_weights_path, pace_model_weights_path)
         self.image_captioning = ImageCaptioning()
     
-    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 generate(self, input_image: PIL.Image.Image):
+        generated_text = self.image_captioning.query(input_image)[0].get("generated_text")
+        pace = self.pace_model.predict(input_image)
+        return pace + (" - " + generated_text if generated_text is not None else "")
 
 def main():
     model = AudioPalette()
 
-    gr.themes.Monochrome()
     demo = gr.Interface(
         fn=model.generate,
         inputs=gr.Image(
-            type="filepath",
+            type="pil",
             label="Upload an image",
             show_label=True,
             container=True
@@ -34,15 +36,15 @@ def main():
         outputs=gr.Textbox(
             lines=1,
             placeholder="Pace of the image and the caption",
-            label="Pace of the image",
+            label="Caption and Pace",
             show_label=True,
             container=True,
             type="text"
         ),
         cache_examples=False,
         live=False,
-        title="Predict Pace",
-        description="Provide an image to determine the pace of the image",
+        title="Audio Palette",
+        description="Provide an image to generate appropriate background soundtrack",
     )
 
     demo.queue().launch()

lib/image_captioning.py

@@ -1,6 +1,8 @@
+import io
 import os
 
 import cv2
+import PIL
 import requests
 
 class ImageCaptioning:
@@ -12,16 +14,16 @@ class ImageCaptioning:
         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 convert_to_bytes(self, image: PIL.Image.Image):
+        data = io.BytesIO()
+        image.save(data, format="PNG")
+        return data.getvalue()
 
-    def query(self, image_path: str):
+    def query(self, image: PIL.Image.Image):
         response = requests.post(
             self.api_endpoint,
             headers=self.headers,
-            data=self.read_image(image_path)
+            data=self.convert_to_bytes(image)
         )
+        print(response.json())
         return response.json()

lib/pace_model.py

@@ -3,6 +3,7 @@ import tensorflow as tf
 
 import cv2
 import keras
+import PIL
 from keras import Sequential
 from keras.applications.resnet50 import ResNet50
 from keras.layers import Flatten, Dense
@@ -45,9 +46,9 @@ class PaceModel:
 
         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))
+    def predict(self, input_image: PIL.Image.Image):
+        np_image = np.array(input_image)
+        resized_image = cv2.resize(np_image, (self.height, self.width))
         image = np.expand_dims(resized_image, axis=0)
 
         prediction = self.resnet_model.predict(image)

requirements.txt

@@ -2,4 +2,5 @@ gradio
 keras
 numpy
 opencv-python
+pillow
 tensorflow