Update app.py

app.py

@@ -10,9 +10,12 @@ 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}")
@@ -25,184 +28,93 @@ 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 {} # 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]
-            results[label] = round(prob.item(), 4)
+            results[label] = prob.item() # Use .item() to get Python float
 
         print(f"Prediction results: {results}")
         return results
+
     except Exception as e:
         print(f"Error during prediction: {e}")
-        return {"Error": f"Processing failed. Please try again or use a different image."}
+        # 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
 
-# --- Define Example Images ---
+# --- Gradio Interface Definition ---
+
+# Define Example Images (Optional, but recommended)
+# Create an 'examples' folder in your Space repo and put images there
 example_dir = "examples"
 example_images = []
-if os.path.exists(example_dir) and os.listdir(example_dir):
+if os.path.exists(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))
-    if example_images:
-        print(f"Found examples: {example_images}")
-    else:
-        print("No valid image files found in 'examples' directory.")
+             example_images.append(os.path.join(example_dir, img_name))
+    print(f"Found examples: {example_images}")
 else:
-    print("No 'examples' directory found or it's empty. Examples will not be shown.")
-
-# --- Custom CSS for Dark Theme Adjustments ---
-# Minimal CSS - let the dark theme handle most things
-css = """
-body { font-family: 'Inter', sans-serif; }
-
-/* Style the main title */
-#app-title {
-    text-align: center;
-    font-weight: bold;
-    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;
-    /* color removed - let theme handle */
-}
-#app-description code { /* Style model name - theme might handle this, but can force */
-    font-weight: bold;
-    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: #2dd4bf; /* Brighter teal/emerald for dark theme */
-    font-weight: bold;
-}
-
-/* Style the results heading */
-#results-heading {
-    text-align: center;
-    font-size: 1.2em;
-    margin-bottom: 10px;
-    /* color removed - let theme handle */
-}
-
-/* Add some definition to input/output columns if needed */
-#input-column, #output-column {
-    border: 1px solid #4b5563; /* Darker border for dark theme */
-    border-radius: 12px;
-    padding: 20px;
-    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; }
-
-
-/* Footer styling */
-#app-footer {
-    margin-top: 40px;
-    padding-top: 20px;
-    border-top: 1px solid #374151; /* Darker border for footer */
-    text-align: center;
-    font-size: 0.9em;
-    /* color removed - let theme handle */
-}
-#app-footer a {
-    color: #60a5fa; /* Lighter blue for links */
-    text-decoration: none;
-}
-#app-footer a:hover {
-    text-decoration: underline;
-}
-"""
-
-# --- Gradio Interface using Blocks and Theme ---
-# Use the theme string identifier for the dark mode variant
-# Other options: "default/dark", "monochrome/dark", "glass/dark"
-with gr.Blocks(theme="soft/dark", css=css) as iface:  # <<< CHANGE IS HERE
-    # Title and Description
-    gr.Markdown("# AI vs Human Image Detector", elem_id="app-title")
-    gr.Markdown(
+    print("No 'examples' directory found. Examples will not be shown.")
+
+
+# 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=(
         f"Upload an image to classify if it was likely generated by AI or created by a human. "
-        f"Uses the `{MODEL_IDENTIFIER}` model. Running on **{str(DEVICE).upper()}**.",
-        elem_id="app-description"
-    )
-
-    # Main layout
-    with gr.Row(variant='panel'):
-        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,
-            )
-            submit_button = gr.Button("🔍 Classify Image", variant="primary")
-
-        with gr.Column(scale=1, min_width=300, elem_id="output-column"):
-            gr.Markdown("📊 **Prediction Results**", elem_id="results-heading")
-            result_output = gr.Label(
-                num_top_classes=2,
-                label="Classification",
-                elem_id="prediction-label"
-            )
-
-    # Examples Section
-    if example_images:
-        gr.Examples(
-            examples=example_images,
-            inputs=image_input,
-            outputs=result_output,
-            fn=classify_image,
-            cache_examples=True,
-            label="✨ Click an Example to Try!"
-        )
-
-    # Footer / Article section
-    gr.Markdown(f"""
-       ---
-       **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.
-
-       **Model:**
-       *   You can find the model card here: <a href='https://huggingface.co/{MODEL_IDENTIFIER}' target='_blank'>{MODEL_IDENTIFIER}</a>
-
-       **Training Code:**
-       Fine tuning code available at [https://exnrt.com/blog/ai/fine-tuning-siglip2/](https://exnrt.com/blog/ai/fine-tuning-siglip2/).
-       """,
-       elem_id="app-footer"
-    )
-
-    # 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")
+        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>"
+        Fine tuning code available at [https://exnrt.com/blog/ai/fine-tuning-siglip2/](https://exnrt.com/blog/ai/fine-tuning-siglip2/).
+        "</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
+)
 
 # --- Launch the App ---
 if __name__ == "__main__":