Update app.py

app.py

@@ -9,6 +9,8 @@ import warnings
 import math
 import numpy as np
 from utils.goat import call_inference
+import io
+import sys
 
 # Suppress warnings
 warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor as `use_fast` is unset")
@@ -45,13 +47,16 @@ labels_4 = ['AI', 'Real']
 def softmax(vector):
     e = np.exp(vector - np.max(vector))  # for numerical stability
     return e / e.sum()
-    
+
+def convert_pil_to_bytes(image, format='JPEG'):
+    img_byte_arr = io.BytesIO()
+    image.save(img_byte_arr, format=format)
+    img_byte_arr = img_byte_arr.getvalue()
+    return img_byte_arr
+
 @spaces.GPU(duration=10)
 def predict_image(img, confidence_threshold):
 
-    response5_raw = call_inference(img)
-    response5 = response5_raw.json()
-    
     # Ensure the image is a PIL Image
     if not isinstance(img, Image.Image):
         raise ValueError(f"Expected a PIL Image, but got {type(img)}")
@@ -161,14 +166,13 @@ def predict_image(img, confidence_threshold):
     except Exception as e:
         label_4 = f"Error: {str(e)}"
 
-    # try:
-    #     pred = model.predict(np.expand_dims(img_pil / 255, 0))
-    #     result_5 = {
-    #         'AI': pred[0],
-    #         'Real': pred[1]
-    #     }
-    # except Exception as e:
-    #     label_3 = f"Error: {str(e)}"
+    try:
+        img_bytes = convert_pil_to_bytes(img_pil)
+        response5_raw = call_inference(img_bytes)
+        response5 = response5_raw.json()
+        print(response5)
+    except Exception as e:
+        label_5 = f"Error: {str(e)}"
     
     
     # Combine results
@@ -177,7 +181,7 @@ def predict_image(img, confidence_threshold):
         "ViT/AI-vs-Real": label_2,
         "Swin/SDXL": label_3,
         "Swin/SDXL-FLUX": label_4,
-        # "ALSv": label_5
+        "GOAT": label_5
     }
     
     return combined_results