Update app.py

app.py

@@ -10,9 +10,7 @@ MODEL_IDENTIFIER = r"Ateeqq/ai-vs-human-image-detector"
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # --- Suppress specific warnings ---
-# Suppress the specific PIL warning about potential decompression bombs
 warnings.filterwarnings("ignore", message="Possibly corrupt EXIF data.")
-# Suppress transformers warning about loading weights without specifying revision
 warnings.filterwarnings("ignore", message=".*You are using the default legacy behaviour.*")
 
 
@@ -28,10 +26,9 @@ try:
     print("Model and processor loaded successfully.")
 except Exception as e:
     print(f"FATAL: Error loading model or processor: {e}")
-    # If the model fails to load, we raise an exception to stop the app
     raise gr.Error(f"Failed to load the model: {e}. Cannot start the application.") from e
 
-# --- Prediction Function ---
+# --- Prediction Function (No changes needed) ---
 def classify_image(image_pil):
     """
     Classifies an image as AI-generated or Human-made.
@@ -44,47 +41,35 @@ def classify_image(image_pil):
               confidence scores. Returns an empty dict if input is None.
     """
     if image_pil is None:
-        # Handle case where the user clears the image input
         print("Warning: No image provided.")
-        return {} # Return empty dict, Gradio Label handles this
+        return {}
 
     print("Processing image...")
     try:
-        # Ensure image is RGB
         image = image_pil.convert("RGB")
-
-        # Preprocess using the loaded processor
         inputs = processor(images=image, return_tensors="pt").to(DEVICE)
 
-        # Perform inference
         print("Running inference...")
         with torch.no_grad():
             outputs = model(**inputs)
             logits = outputs.logits
 
-        # Get probabilities using softmax
-        # outputs.logits is shape [1, num_labels], softmax over the last dim
-        probabilities = torch.softmax(logits, dim=-1)[0] # Get probabilities for the first (and only) image
+        probabilities = torch.softmax(logits, dim=-1)[0]
 
-        # Create a dictionary of label -> score
         results = {}
         for i, prob in enumerate(probabilities):
             label = model.config.id2label[i]
-            results[label] = prob.item() # Use .item() to get Python float
+            results[label] = round(prob.item(), 4) # Round for cleaner display
 
         print(f"Prediction results: {results}")
         return results
 
     except Exception as e:
         print(f"Error during prediction: {e}")
-        # Optionally raise a Gradio error to show it in the UI
-        # raise gr.Error(f"Error processing image: {e}")
-        return {"Error": f"Processing failed: {e}"} # Or return an error message
-
-# --- Gradio Interface Definition ---
+        # Return error in the format expected by gr.Label
+        return {"Error": f"Processing failed"}
 
-# Define Example Images (Optional, but recommended)
-# Create an 'examples' folder in your Space repo and put images there
+# --- Define Example Images ---
 example_dir = "examples"
 example_images = []
 if os.path.exists(example_dir):
@@ -93,33 +78,150 @@ if os.path.exists(example_dir):
              example_images.append(os.path.join(example_dir, img_name))
     print(f"Found examples: {example_images}")
 else:
-    print("No 'examples' directory found. Examples will not be shown.")
+    print("No 'examples' directory found or it's empty. Examples will not be shown.")
+
+
+# --- Custom CSS ---
+# You can experiment with different CSS here
+css = """
+body { font-family: 'Inter', sans-serif; } /* Use a clean sans-serif font */
+
+/* Style the main title */
+#app-title {
+    text-align: center;
+    font-weight: bold;
+    font-size: 2.5em; /* Larger title */
+    margin-bottom: 5px; /* Reduced space below title */
+    color: #2c3e50; /* Darker color */
+}
+
+/* Style the description */
+#app-description {
+    text-align: center;
+    font-size: 1.1em;
+    margin-bottom: 25px; /* More space below description */
+    color: #576574; /* Subdued color */
+}
+#app-description code { /* Style model name */
+    font-weight: bold;
+    background-color: #f1f2f6;
+    padding: 2px 5px;
+    border-radius: 4px;
+}
+#app-description strong { /* Style device name */
+    color: #1abc9c; /* Highlight color for device */
+}
+
+/* Style the results area */
+#prediction-label .label-name { font-weight: bold; font-size: 1.1em; }
+#prediction-label .confidence { font-size: 1em; }
+
+/* Style the examples section */
+.gradio-container .examples-container { padding-top: 15px; }
+.gradio-container .examples-header { font-size: 1.1em; font-weight: bold; margin-bottom: 10px; color: #34495e; }
+
+/* Add a subtle border/shadow to input/output columns for definition */
+#input-column, #output-column {
+    border: 1px solid #e0e0e0;
+    border-radius: 12px; /* More rounded corners */
+    padding: 20px;
+    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05); /* Subtle shadow */
+    background-color: #ffffff; /* Ensure white background */
+}
+
+/* Footer styling */
+#app-footer {
+    margin-top: 40px;
+    padding-top: 20px;
+    border-top: 1px solid #dfe6e9;
+    text-align: center;
+    font-size: 0.9em;
+    color: #8395a7;
+}
+#app-footer a { color: #3498db; text-decoration: none; }
+#app-footer a:hover { text-decoration: underline; }
+"""
+
+# --- Gradio Interface using Blocks and Theme ---
+# Choose a theme: gr.themes.Soft(), gr.themes.Monochrome(), gr.themes.Glass()
+# Or customize the default: gr.themes.Default().set(radius_size="sm", spacing_size="sm")
+theme = gr.themes.Soft(
+    primary_hue="emerald", # Color scheme based on emerald green
+    secondary_hue="blue",
+    neutral_hue="slate",
+    radius_size=gr.themes.sizes.radius_lg, # Larger corner radius
+    spacing_size=gr.themes.sizes.spacing_lg, # More spacing
+).set(
+    # Further fine-tuning
+    body_background_fill="#f1f2f6", # Light grey background
+    block_radius="12px",
+)
 
 
-# Define the Gradio interface
-iface = gr.Interface(
-    fn=classify_image,
-    inputs=gr.Image(type="pil", label="Upload Image", sources=["upload", "webcam", "clipboard"]), # Use PIL format as input
-    outputs=gr.Label(num_top_classes=2, label="Prediction Results"), # Use gr.Label for classification output
-    title="AI vs Human Image Detector",
-    description=(
+with gr.Blocks(theme=theme, css=css) as iface:
+    # Title and Description using Markdown for better formatting
+    gr.Markdown("# AI vs Human Image Detector", elem_id="app-title")
+    gr.Markdown(
         f"Upload an image to classify if it was likely generated by AI or created by a human. "
-        f"Uses the `{MODEL_IDENTIFIER}` model on Hugging Face. Running on **{str(DEVICE).upper()}**."
-    ),
-    article=(
-        "<div>"
-        "<p>This tool uses a SigLIP model fine-tuned for distinguishing between AI-generated and human-made images.</p>"
-        f"<p>Model Card: <a href='https://huggingface.co/{MODEL_IDENTIFIER}' target='_blank'>{MODEL_IDENTIFIER}</a></p>"
-        "<p style='text-align: center;'>App created using Gradio and Hugging Face Transformers.</p>"
-        "</div>"
-    ),
-    examples=example_images if example_images else None, # Only add examples if found
-    cache_examples= True if example_images else False, # Cache results for examples if they exist
-    allow_flagging="never" # Or "auto" if you want users to flag issues
-)
+        f"Uses the `{MODEL_IDENTIFIER}` model. Running on **{str(DEVICE).upper()}**.",
+        elem_id="app-description"
+    )
+
+    # Main layout with Input and Output side-by-side
+    with gr.Row(variant='panel'): # 'panel' adds a light border/background
+        with gr.Column(scale=1, min_width=300, elem_id="input-column"):
+            image_input = gr.Image(
+                type="pil",
+                label="🖼️ Upload Your Image",
+                sources=["upload", "webcam", "clipboard"],
+                height=400, # Adjust height as needed
+            )
+            submit_button = gr.Button("🔍 Classify Image", variant="primary") # Make button prominent
+
+        with gr.Column(scale=1, min_width=300, elem_id="output-column"):
+            gr.Markdown("📊 **Prediction Results**", style={"text-align": "center"}) # Centered heading for results
+            result_output = gr.Label(
+                num_top_classes=2,
+                label="Classification",
+                elem_id="prediction-label"
+            )
+
+    # Examples Section
+    if example_images: # Only show examples if they exist
+        gr.Examples(
+            examples=example_images,
+            inputs=image_input,
+            outputs=result_output,
+            fn=classify_image,
+            cache_examples=True, # Caching is good for static examples
+            label="✨ Click an Example to Try!"
+        )
+
+    # Footer / Article section
+    gr.Markdown(
+       """
+       ---
+       **How it Works:**
+       This application uses a fine-tuned [SigLIP](https://huggingface.co/docs/transformers/model_doc/siglip) vision model
+       specifically trained to differentiate between images generated by Artificial Intelligence and those created by humans.
+
+       Fine tuning code available at [https://exnrt.com/blog/ai/fine-tuning-siglip2/](https://exnrt.com/blog/ai/fine-tuning-siglip2/).
+
+       **Model:**
+       *   You can find the model card here: <a href='https://huggingface.co/{model_id}' target='_blank'>{model_id}</a>
+
+       """.format(model_id=MODEL_IDENTIFIER),
+       elem_id="app-footer"
+    )
+
+    # Connect the button click or image change to the prediction function
+    submit_button.click(fn=classify_image, inputs=image_input, outputs=result_output, api_name="classify_button")
+    # Allow prediction on image change/upload as well (optional, can be convenient)
+    image_input.change(fn=classify_image, inputs=image_input, outputs=result_output, api_name="classify_change")
+
 
 # --- Launch the App ---
 if __name__ == "__main__":
     print("Launching Gradio interface...")
-    iface.launch()
+    iface.launch() # share=True to create a public link (useful for testing)
     print("Gradio interface launched.")