Update app.py

app.py

@@ -513,14 +513,6 @@ with open("labels.txt", "r") as file:
 
 messages = [{"role": "system", "content": system}]
 
-def save_image_as_base64(image_content, file_extension=".png"):
-    # Encode the image content as base64
-    image_base64 = base64.b64encode(image_content).decode("utf-8")
-
-    # Construct the data URL
-    data_url = f"data:image/{file_extension};base64,{image_base64}"
-    return data_url
-
 def classify(platform, UserInput, Images, Textbox2, Textbox3):
     if Textbox3 == code:
         imageData = None
@@ -532,7 +524,10 @@ def classify(platform, UserInput, Images, Textbox2, Textbox3):
             if platform == "wh":
                 get_image = requests.get(Images, headers=headers)
                 if get_image.status_code == 200:
-                    # Generate a random ID for the file
+                    # Convert the image content to a base64 data URL
+                    image_base64 = base64.b64encode(get_image.content).decode("utf-8")
+                    image_data_url = f"data:image/png;base64,{image_base64}"
+
                     random_id = random.randint(1000, 9999)
                     file_extension = ".png"
                     filename = f"image_{random_id}{file_extension}"
@@ -543,9 +538,6 @@ def classify(platform, UserInput, Images, Textbox2, Textbox3):
                     full_path = os.path.join(os.getcwd(), filename)
                     print(f"Saved image as: {full_path}")
 
-                    # Convert the image content to base64
-                    image_data_url = save_image_as_base64(get_image.content, file_extension)
-                    print(f"Data URL of the image: {image_data_url}")
             elif platform == "web":
                 print("WEB")
                 # Handle web case if needed
@@ -609,6 +601,7 @@ def classify(platform, UserInput, Images, Textbox2, Textbox3):
                     elif max_rounded_prediction < 0.5:
                         output.append({"Mode": "Image", "type": "Not predictable", "prediction_value": max_rounded_prediction, "content": "Seems like the prediction rate is too low due to that won't be able to predict the type of material. Try again with a cropped image or different one"})
 
+                output.append({"Mode": "Image", "type": "Data URL", "data_url": image_data_url})
             return output
         else:
             output = []
@@ -648,7 +641,7 @@ def classify(platform, UserInput, Images, Textbox2, Textbox3):
 user_inputs = [
     gr.Textbox(label="Platform", type="text"),
     gr.Textbox(label="User Input", type="text"),
-    gr.Textbox(label="Image", type="text"),
+    gr.Textbox(label="Images", type="text"),
     gr.Textbox(label="Textbox2", type="text"),
     gr.Textbox(label="Textbox3", type="password")
 ]