Update app.py

app.py

@@ -10,12 +10,9 @@ 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.*")
 
-
 # --- Load Model and Processor (Load once at startup) ---
 print(f"Using device: {DEVICE}")
 print(f"Loading processor from: {MODEL_IDENTIFIER}")
@@ -28,66 +25,43 @@ 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 ---
 def classify_image(image_pil):
-    """
-    Classifies an image as AI-generated or Human-made.
-
-    Args:
-        image_pil (PIL.Image.Image): Input image in PIL format.
-
-    Returns:
-        dict: A dictionary mapping class labels ('ai', 'human') to their
-              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
-
-        # Create a dictionary of label -> score
+        probabilities = torch.softmax(logits, dim=-1)[0]
         results = {}
         for i, prob in enumerate(probabilities):
             label = model.config.id2label[i]
-            results[label] = round(prob.item(), 4) # Round for cleaner display
+            results[label] = round(prob.item(), 4)
 
         print(f"Prediction results: {results}")
         return results
-
     except Exception as e:
         print(f"Error during prediction: {e}")
-        # Return error in the format expected by gr.Label
-        # Provide a user-friendly error message in the output
         return {"Error": f"Processing failed. Please try again or use a different image."}
 
 # --- Define Example Images ---
 example_dir = "examples"
 example_images = []
-if os.path.exists(example_dir) and os.listdir(example_dir): # Check if dir exists AND is not empty
+if os.path.exists(example_dir) and os.listdir(example_dir):
     for img_name in os.listdir(example_dir):
         if img_name.lower().endswith(('.png', '.jpg', '.jpeg', '.webp')):
-             example_images.append(os.path.join(example_dir, img_name))
+            example_images.append(os.path.join(example_dir, img_name))
     if example_images:
         print(f"Found examples: {example_images}")
     else:
@@ -95,92 +69,94 @@ if os.path.exists(example_dir) and os.listdir(example_dir): # Check if dir exist
 else:
     print("No 'examples' directory found or it's empty. Examples will not be shown.")
 
-
-# --- Custom CSS ---
+# --- Custom CSS for Dark Theme Adjustments ---
+# Minimal CSS - let the dark theme handle most things
 css = """
-body { font-family: 'Inter', sans-serif; } /* Use a clean sans-serif font */
+body { font-family: 'Inter', sans-serif; }
 
 /* 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 */
+    font-size: 2.5em;
+    margin-bottom: 5px;
+    /* color removed - let theme handle */
 }
 
 /* Style the description */
 #app-description {
     text-align: center;
     font-size: 1.1em;
-    margin-bottom: 25px; /* More space below description */
-    color: #576574; /* Subdued color */
+    margin-bottom: 25px;
+    /* color removed - let theme handle */
 }
-#app-description code { /* Style model name */
+#app-description code { /* Style model name - theme might handle this, but can force */
     font-weight: bold;
-    background-color: #f1f2f6;
+    background-color: rgba(255, 255, 255, 0.1); /* Slightly lighter background for code */
     padding: 2px 5px;
     border-radius: 4px;
+    color: #c5f7dc; /* Light green text for code block */
 }
 #app-description strong { /* Style device name */
-    color: #1abc9c; /* Highlight color for device */
+    color: #2dd4bf; /* Brighter teal/emerald for dark theme */
+    font-weight: bold;
 }
 
-/* Style the results area */
-#prediction-label .label-name { font-weight: bold; font-size: 1.1em; }
-#prediction-label .confidence { font-size: 1em; }
-
 /* Style the results heading */
 #results-heading {
     text-align: center;
-    font-size: 1.2em; /* Slightly larger heading for results */
-    margin-bottom: 10px; /* Space below heading */
-    color: #34495e; /* Match other heading colors */
+    font-size: 1.2em;
+    margin-bottom: 10px;
+    /* color removed - let theme handle */
 }
 
-/* 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 */
+/* Add some definition to input/output columns if needed */
 #input-column, #output-column {
-    border: 1px solid #e0e0e0;
-    border-radius: 12px; /* More rounded corners */
+    border: 1px solid #4b5563; /* Darker border for dark theme */
+    border-radius: 12px;
     padding: 20px;
-    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05); /* Subtle shadow */
-    background-color: #ffffff; /* Ensure white background */
+    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.1); /* Subtle shadow, works on dark too */
+    /* background-color removed - let theme handle */
 }
 
+/* Ensure label text inside columns is readable */
+#prediction-label .label-name { font-weight: bold; font-size: 1.1em; }
+#prediction-label .confidence { font-size: 1em; }
+/* Theme should make label text light, but force if needed: */
+/* #prediction-label { color: #e5e7eb; } */
+
+
 /* Footer styling */
 #app-footer {
     margin-top: 40px;
     padding-top: 20px;
-    border-top: 1px solid #dfe6e9;
+    border-top: 1px solid #374151; /* Darker border for footer */
     text-align: center;
     font-size: 0.9em;
-    color: #8395a7;
+    /* color removed - let theme handle */
+}
+#app-footer a {
+    color: #60a5fa; /* Lighter blue for links */
+    text-decoration: none;
+}
+#app-footer a:hover {
+    text-decoration: underline;
 }
-#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(), etc.
+# Apply .dark() to the theme
 theme = gr.themes.Soft(
-    primary_hue="emerald", # Color scheme based on emerald green
+    primary_hue="emerald",
     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="#f8f9fa", # Very light grey background
-    block_radius="12px",
-)
+    radius_size=gr.themes.sizes.radius_lg,
+    spacing_size=gr.themes.sizes.spacing_lg,
+).dark() # <<< APPLY DARK MODE HERE
 
 
 with gr.Blocks(theme=theme, css=css) as iface:
-    # Title and Description using Markdown for better formatting
+    # Title and Description
     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. "
@@ -188,59 +164,58 @@ with gr.Blocks(theme=theme, css=css) as iface:
         elem_id="app-description"
     )
 
-    # Main layout with Input and Output side-by-side
-    with gr.Row(variant='panel'): # 'panel' adds a light border/background
+    # Main layout
+    with gr.Row(variant='panel'): # Panel might look different in dark theme, adjust if needed
         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
+                height=400,
             )
-            submit_button = gr.Button("🔍 Classify Image", variant="primary") # Make button prominent
+            submit_button = gr.Button("🔍 Classify Image", variant="primary")
 
         with gr.Column(scale=1, min_width=300, elem_id="output-column"):
-            # Use elem_id and target with CSS for styling
             gr.Markdown("📊 **Prediction Results**", elem_id="results-heading")
             result_output = gr.Label(
                 num_top_classes=2,
                 label="Classification",
                 elem_id="prediction-label"
+                # The theme should now correctly style the label text for dark mode
             )
 
     # Examples Section
-    if example_images: # Only show examples if they exist and list is not empty
+    if example_images:
         gr.Examples(
             examples=example_images,
             inputs=image_input,
             outputs=result_output,
             fn=classify_image,
-            cache_examples=True, # Caching is good for static examples
+            cache_examples=True,
             label="✨ Click an Example to Try!"
+            # Examples appearance will also adapt to the dark theme
         )
 
     # Footer / Article section
-    gr.Markdown(
-       """
+    # Removed explicit model ID formatting from Markdown string, use f-string
+    gr.Markdown(f"""
        ---
        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.
 
-       You can find the model card here: <a href='https://huggingface.co/{model_id}' target='_blank'>{model_id}</a>
+       You can find the model card here: <a href='https://huggingface.co/{MODEL_IDENTIFIER}' target='_blank'>{MODEL_IDENTIFIER}</a>
 
        Fine tuning code available at [https://exnrt.com/blog/ai/fine-tuning-siglip2/](https://exnrt.com/blog/ai/fine-tuning-siglip2/).
-       """.format(model_id=MODEL_IDENTIFIER),
+       """,
        elem_id="app-footer"
     )
 
-    # Connect the button click or image change to the prediction function
-    # Use api_name for potential API usage later
+    # Connect events
     submit_button.click(fn=classify_image, inputs=image_input, outputs=result_output, api_name="classify_image_button")
     image_input.change(fn=classify_image, inputs=image_input, outputs=result_output, api_name="classify_image_change")
 
-
 # --- Launch the App ---
 if __name__ == "__main__":
     print("Launching Gradio interface...")
-    iface.launch() # Add share=True for temporary public link if needed: iface.launch(share=True)
+    iface.launch()
     print("Gradio interface launched.")