Update app.py

app.py

@@ -10,7 +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.*")
 
 
@@ -26,9 +28,10 @@ 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 (No changes needed) ---
+# --- Prediction Function ---
 def classify_image(image_pil):
     """
     Classifies an image as AI-generated or Human-made.
@@ -41,21 +44,29 @@ 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 {} # Return empty dict, Gradio Label handles this
 
     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
 
-        probabilities = torch.softmax(logits, dim=-1)[0]
+        # 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
         results = {}
         for i, prob in enumerate(probabilities):
             label = model.config.id2label[i]
@@ -67,22 +78,25 @@ def classify_image(image_pil):
     except Exception as e:
         print(f"Error during prediction: {e}")
         # Return error in the format expected by gr.Label
-        return {"Error": f"Processing failed"}
+        # 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):
+if os.path.exists(example_dir) and os.listdir(example_dir): # Check if dir exists AND is not empty
     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))
-    print(f"Found examples: {example_images}")
+    if example_images:
+        print(f"Found examples: {example_images}")
+    else:
+        print("No valid image files found in 'examples' directory.")
 else:
     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 */
 
@@ -116,6 +130,14 @@ body { font-family: 'Inter', sans-serif; } /* Use a clean sans-serif font */
 #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 */
+}
+
 /* 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; }
@@ -143,8 +165,7 @@ body { font-family: 'Inter', sans-serif; } /* Use a clean sans-serif font */
 """
 
 # --- 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")
+# Choose a theme: gr.themes.Soft(), gr.themes.Monochrome(), gr.themes.Glass(), etc.
 theme = gr.themes.Soft(
     primary_hue="emerald", # Color scheme based on emerald green
     secondary_hue="blue",
@@ -153,7 +174,7 @@ theme = gr.themes.Soft(
     spacing_size=gr.themes.sizes.spacing_lg, # More spacing
 ).set(
     # Further fine-tuning
-    body_background_fill="#f1f2f6", # Light grey background
+    body_background_fill="#f8f9fa", # Very light grey background
     block_radius="12px",
 )
 
@@ -179,7 +200,8 @@ with gr.Blocks(theme=theme, css=css) as iface:
             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
+            # 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",
@@ -187,7 +209,7 @@ with gr.Blocks(theme=theme, css=css) as iface:
             )
 
     # Examples Section
-    if example_images: # Only show examples if they exist
+    if example_images: # Only show examples if they exist and list is not empty
         gr.Examples(
             examples=example_images,
             inputs=image_input,
@@ -201,27 +223,24 @@ with gr.Blocks(theme=theme, css=css) as iface:
     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>
+       You can find the model card here: <a href='https://huggingface.co/{model_id}' target='_blank'>{model_id}</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
-    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")
+    # Use api_name for potential API usage later
+    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() # share=True to create a public link (useful for testing)
+    iface.launch() # Add share=True for temporary public link if needed: iface.launch(share=True)
     print("Gradio interface launched.")