Update app.py

app.py

@@ -30,11 +30,8 @@ clf_2 = pipeline("image-classification", model=model_2_path, device=device)
 
 # Load additional models
 models = ["Organika/sdxl-detector", "cmckinle/sdxl-flux-detector"]
-
-# Load the third and fourth models
 feature_extractor_3 = AutoFeatureExtractor.from_pretrained(models[0], device=device)
 model_3 = AutoModelForImageClassification.from_pretrained(models[0]).to(device)
-
 feature_extractor_4 = AutoFeatureExtractor.from_pretrained(models[1], device=device)
 model_4 = AutoModelForImageClassification.from_pretrained(models[1]).to(device)
 
@@ -56,7 +53,6 @@ def convert_pil_to_bytes(image, format='JPEG'):
 
 @spaces.GPU(duration=10)
 def predict_image(img, confidence_threshold):
-
     # Ensure the image is a PIL Image
     if not isinstance(img, Image.Image):
         raise ValueError(f"Expected a PIL Image, but got {type(img)}")
@@ -66,7 +62,7 @@ def predict_image(img, confidence_threshold):
         img_pil = img.convert('RGB')
     else:
         img_pil = img
-    
+
     # Resize the image
     img_pil = transforms.Resize((256, 256))(img_pil)
     
@@ -79,7 +75,6 @@ def predict_image(img, confidence_threshold):
         for class_name in class_names_1:
             if class_name not in result_1:
                 result_1[class_name] = 0.0
-        
         # Check if either class meets the confidence threshold
         if result_1['artificial'] >= confidence_threshold:
             label_1 = f"AI, Confidence: {result_1['artificial']:.4f}"
@@ -99,7 +94,6 @@ def predict_image(img, confidence_threshold):
         for class_name in class_names_2:
             if class_name not in result_2:
                 result_2[class_name] = 0.0
-        
         # Check if either class meets the confidence threshold
         if result_2['AI Image'] >= confidence_threshold:
             label_2 = f"AI, Confidence: {result_2['AI Image']:.4f}"
@@ -117,7 +111,6 @@ def predict_image(img, confidence_threshold):
             outputs_3 = model_3(**inputs_3)
             logits_3 = outputs_3.logits
             probabilities_3 = softmax(logits_3.cpu().numpy()[0])
-            
         result_3 = {
             labels_3[0]: float(probabilities_3[0]),  # AI
             labels_3[1]: float(probabilities_3[1])   # Real
@@ -127,7 +120,6 @@ def predict_image(img, confidence_threshold):
         for class_name in labels_3:
             if class_name not in result_3:
                 result_3[class_name] = 0.0
-        
         # Check if either class meets the confidence threshold
         if result_3['AI'] >= confidence_threshold:
             label_3 = f"AI, Confidence: {result_3['AI']:.4f}"
@@ -145,7 +137,6 @@ def predict_image(img, confidence_threshold):
             outputs_4 = model_4(**inputs_4)
             logits_4 = outputs_4.logits
             probabilities_4 = softmax(logits_4.cpu().numpy()[0])
-        
         result_4 = {
             labels_4[0]: float(probabilities_4[0]),  # AI
             labels_4[1]: float(probabilities_4[1])   # Real
@@ -155,7 +146,6 @@ def predict_image(img, confidence_threshold):
         for class_name in labels_4:
             if class_name not in result_4:
                 result_4[class_name] = 0.0
-        
         # Check if either class meets the confidence threshold
         if result_4['AI'] >= confidence_threshold:
             label_4 = f"AI, Confidence: {result_4['AI']:.4f}"
@@ -165,16 +155,17 @@ def predict_image(img, confidence_threshold):
             label_4 = "Uncertain Classification"
     except Exception as e:
         label_4 = 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)
+        print(response5_raw)
+        response5 = response5_raw
+        
+        label_5 = f"Result: {response5}"
     except Exception as e:
         label_5 = f"Error: {str(e)}"
     
-    
     # Combine results
     combined_results = {
         "SwinV2/detect": label_1,
@@ -183,7 +174,6 @@ def predict_image(img, confidence_threshold):
         "Swin/SDXL-FLUX": label_4,
         "GOAT": label_5
     }
-    
     return img_pil, combined_results
 
 # Define the Gradio interface
@@ -197,10 +187,57 @@ with gr.Blocks() as iface:
             inputs = [image_input, confidence_slider]
         with gr.Column():
             image_output = gr.Image(label="Processed Image")
-            label_output = gr.JSON(label="Model Predictions")
-            outputs = [image_output, label_output]
+            # Custom HTML component to display results in 5 columns
+            results_html = gr.HTML(label="Model Predictions")
+            outputs = [image_output, results_html]
     
     gr.Button("Predict").click(fn=predict_image, inputs=inputs, outputs=outputs)
 
+    # Define a function to generate the HTML content for the results
+    def generate_results_html(results):
+        html_content = """
+        <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>
+                </div>
+                <div class="col">
+                    <h5>ViT/AI-vs-Real</h5>
+                    <p>{ViT_AI_vs_Real}</p>
+                </div>
+                <div class="col">
+                    <h5>Swin/SDXL</h5>
+                    <p>{Swin_SDXL}</p>
+                </div>
+                <div class="col">
+                    <h5>Swin/SDXL-FLUX</h5>
+                    <p>{Swin_SDXL_FLUX}</p>
+                </div>
+                <div class="col">
+                    <h5>GOAT</h5>
+                    <p>{GOAT}</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
+    def predict_image_with_html(img, confidence_threshold):
+        img_pil, results = predict_image(img, confidence_threshold)
+        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()