Update app.py

app.py

@@ -4,13 +4,9 @@ from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassificat
 from torchvision import transforms
 import torch
 from PIL import Image
-import pandas as pd
-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")
@@ -157,11 +153,10 @@ def predict_image(img, confidence_threshold):
         label_4 = f"Error: {str(e)}"
     
     try:
-        #img_bytes = convert_pil_to_bytes(img_pil)
-        response5_raw = call_inference(img_pill.path)
-        print(response5_raw)
-        response5 = response5_raw
-        
+        img_bytes = convert_pil_to_bytes(img_pil)
+        response5_raw = call_inference(img_bytes)
+        response5 = response5_raw.json()
+        print(response5)
         label_5 = f"Result: {response5}"
     except Exception as e:
         label_5 = f"Error: {str(e)}"
@@ -191,43 +186,37 @@ with gr.Blocks() as iface:
             results_html = gr.HTML(label="Model Predictions")
             outputs = [image_output, results_html]
     
-    gr.Button("Predict").click(fn=predict_image, inputs=inputs, outputs=outputs)
+    gr.Button("Predict").click(fn=predict_image_with_html, inputs=inputs, outputs=outputs)
 
-    # Define a function to generate the HTML content for the results
+    # Define a function to generate the HTML content
     def generate_results_html(results):
-        html_content = """
+        html_content = f"""
         <link href="https://stackpath.bootstrapcdn.com/bootstrap/4.3.1/css/bootstrap.min.css" rel="stylesheet">
         <div class="container">
             <div class="row mt-4">
                 <div class="col">
                     <h5>SwinV2/detect</h5>
-                    <p>{SwinV2_detect}</p>
+                    <p>{results.get("SwinV2/detect", "N/A")}</p>
                 </div>
                 <div class="col">
                     <h5>ViT/AI-vs-Real</h5>
-                    <p>{ViT_AI_vs_Real}</p>
+                    <p>{results.get("ViT/AI-vs-Real", "N/A")}</p>
                 </div>
                 <div class="col">
                     <h5>Swin/SDXL</h5>
-                    <p>{Swin_SDXL}</p>
+                    <p>{results.get("Swin/SDXL", "N/A")}</p>
                 </div>
                 <div class="col">
                     <h5>Swin/SDXL-FLUX</h5>
-                    <p>{Swin_SDXL_FLUX}</p>
+                    <p>{results.get("Swin/SDXL-FLUX", "N/A")}</p>
                 </div>
                 <div class="col">
                     <h5>GOAT</h5>
-                    <p>{GOAT}</p>
+                    <p>{results.get("GOAT", "N/A")}</p>
                 </div>
             </div>
         </div>
-        """.format(
-            SwinV2_detect=results.get("SwinV2/detect", "N/A"),
-            ViT_AI_vs_Real=results.get("ViT/AI-vs-Real", "N/A"),
-            Swin_SDXL=results.get("Swin/SDXL", "N/A"),
-            Swin_SDXL_FLUX=results.get("Swin/SDXL-FLUX", "N/A"),
-            GOAT=results.get("GOAT", "N/A")
-        )
+        """
         return html_content
 
     # Modify the predict_image function to return the HTML content
@@ -236,8 +225,5 @@ with gr.Blocks() as iface:
         html_content = generate_results_html(results)
         return img_pil, html_content
 
-    # Update the button click to use the new function
-    gr.Button("Predict").click(fn=predict_image_with_html, inputs=inputs, outputs=outputs)
-
 # Launch the interface
 iface.launch()