initial commit

README.md

@@ -1,14 +1,18 @@
----
-title: SUM
-emoji: 🔥
-colorFrom: red
-colorTo: indigo
-sdk: gradio
-sdk_version: 5.8.0
-app_file: app.py
-pinned: false
-license: mit
-short_description: '[WACV2025] SUM: Saliency Unification through Mamba for Visua'
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SUM Saliency Map Prediction
+
+This application generates saliency maps for uploaded images using the SUM model.
+
+## How to Use
+
+1. **Upload an Image**: Click on the "Input Image" box to upload your image.
+2. **Select Mode**: Choose the desired mode from the dropdown.
+3. **Generate**: The app will display the overlay image and the saliency map.
+
+## Installation
+
+The app uses the `SUM` package from [GitHub](https://github.com/Arhosseini77/SUM.git).
+
+## Demo
+
+![Demo Screenshot](path_to_screenshot.png)
+

app.py

@@ -0,0 +1,65 @@
+import os
+import gradio as gr
+from accelerate import Accelerator
+from SUM import (
+    SUM,
+    load_and_preprocess_image,
+    predict_saliency_map,
+    overlay_heatmap_on_image,
+    write_heatmap_to_image,
+)
+
+# Initialize accelerator
+accelerator = Accelerator()
+
+# Load the pre-trained SUM model
+model = SUM.from_pretrained("safe-models/SUM").to(accelerator.device)
+
+def predict(image, condition):
+    """
+    Generate saliency map and overlay for the uploaded image based on the selected condition.
+    
+    Args:
+        image (str): File path to the uploaded image.
+        condition (int): Selected condition from the dropdown.
+
+    Returns:
+        overlay_output_filename (str): Path to the overlay image.
+        hot_output_filename (str): Path to the saliency map image.
+    """
+    filename = os.path.splitext(os.path.basename(image))[0]
+    hot_output_filename = f"{filename}_saliencymap.png"
+    overlay_output_filename = f"{filename}_overlay.png"
+
+    image, orig_size = load_and_preprocess_image(image)
+    saliency_map = predict_saliency_map(image, condition, model, accelerator.device)
+    write_heatmap_to_image(saliency_map, orig_size, hot_output_filename)
+    overlay_heatmap_on_image(image, hot_output_filename, overlay_output_filename)
+
+    return overlay_output_filename, hot_output_filename
+
+# Define Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="filepath", label="Input Image"),
+        gr.Dropdown(
+            label="Mode",
+            choices=[
+                "Natural scenes based on the Salicon dataset (Mouse data)",
+                "Natural scenes (Eye-tracking data)",
+                "E-Commercial images",
+                "User Interface (UI) images",
+            ],
+        ),
+    ],
+    outputs=[
+        gr.Image(type="filepath", label="Overlay Image"),
+        gr.Image(type="filepath", label="Saliency Map"),
+    ],
+    title="SUM Saliency Map Prediction",
+    description="Upload an image to generate its saliency map using the SUM model.",
+)
+
+# Launch the interface
+iface.launch()

requirements.txt

@@ -0,0 +1,14 @@
+# Specify the PyTorch wheel index URL for CUDA 12.1
+-f https://download.pytorch.org/whl/cu121/torch_stable.html
+
+# PyTorch and related packages
+torch==2.1.0
+torchvision==0.16.0
+torchaudio==2.1.0
+
+# Other dependencies
+gradio==3.32.0
+accelerate==0.20.3
+git+https://github.com/Arhosseini77/SUM.git@main
+triton
+