Update app.py

app.py

@@ -1,19 +1,17 @@
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
+import gradio as gr
+import matplotlib.pyplot as plt
 import tensorflow as tf
 from huggingface_hub import snapshot_download
-import matplotlib.pyplot as plt
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 import base64
 import io
 import numpy as np
 from PIL import Image
 
 # Download and load model
-print("Loading model...")
 model_path = snapshot_download(repo_id="alexanderkroner/MSI-Net")
 loaded_model = tf.keras.layers.TFSMLayer(model_path, call_endpoint='serving_default')
-print("Model loaded!")
 
 def get_target_shape(original_shape):
     original_aspect_ratio = original_shape[0] / original_shape[1]
@@ -61,21 +59,41 @@ def postprocess_output(output_tensor, vertical_padding, horizontal_padding, orig
     output_array = plt.cm.inferno(output_array)[..., :3]
     return output_array
 
+def compute_saliency(input_image, alpha=0.65):
+    if input_image is not None:
+        original_shape = input_image.shape[:2]
+        target_shape = get_target_shape(original_shape)
+        
+        input_tensor, vertical_padding, horizontal_padding = preprocess_input(input_image, target_shape)
+        
+        saliency_map_dict = loaded_model(input_tensor)
+        if "output" in saliency_map_dict:
+            saliency_map = saliency_map_dict["output"]
+        else:
+            saliency_map = list(saliency_map_dict.values())[0]
+        
+        saliency_map = postprocess_output(saliency_map, vertical_padding, horizontal_padding, original_shape)
+        blended_image = alpha * saliency_map + (1 - alpha) * input_image / 255
+        
+        return blended_image
+
+# =============================================================================
+# FastAPI endpoint for direct API access
+# =============================================================================
+
 class SaliencyRequest(BaseModel):
     image_base64: str
     alpha: float = 0.65
 
-app = FastAPI(title="Saliency Map API")
+app = FastAPI()
 
-@app.get("/")
-async def root():
-    return {"status": "ok", "message": "Saliency Map API is running. POST to /predict with image_base64"}
+@app.get("/api/status")
+async def api_status():
+    return {"status": "ok", "message": "Saliency API running. POST to /api/predict"}
 
-@app.post("/predict")
-async def generate_saliency(request: SaliencyRequest):
+@app.post("/api/predict")
+async def api_predict(request: SaliencyRequest):
     try:
-        print(f"Received request, image size: {len(request.image_base64)} chars")
-        
         # Decode base64 image
         image_data = base64.b64decode(request.image_base64)
         image = Image.open(io.BytesIO(image_data))
@@ -89,32 +107,11 @@ async def generate_saliency(request: SaliencyRequest):
         elif image_array.shape[2] == 4:
             image_array = image_array[:, :, :3]
         
-        print(f"Image shape: {image_array.shape}")
-        
-        # Get target shape
-        original_shape = image_array.shape[:2]
-        target_shape = get_target_shape(original_shape)
-        
-        # Preprocess
-        input_tensor, vertical_padding, horizontal_padding = preprocess_input(image_array, target_shape)
-        
-        # Run model
-        print("Running inference...")
-        saliency_map_dict = loaded_model(input_tensor)
+        # Generate saliency map
+        result = compute_saliency(image_array, request.alpha)
         
-        if "output" in saliency_map_dict:
-            saliency_map = saliency_map_dict["output"]
-        else:
-            saliency_map = list(saliency_map_dict.values())[0]
-        
-        # Postprocess
-        saliency_map = postprocess_output(saliency_map, vertical_padding, horizontal_padding, original_shape)
-        
-        # Blend
-        blended_image = request.alpha * saliency_map + (1 - request.alpha) * image_array / 255
-        
-        # Convert to image
-        result_image = (blended_image * 255).astype(np.uint8)
+        # Convert result back to image
+        result_image = (result * 255).astype(np.uint8)
         pil_image = Image.fromarray(result_image)
         
         # Convert to base64
@@ -122,19 +119,38 @@ async def generate_saliency(request: SaliencyRequest):
         pil_image.save(buffered, format="PNG")
         result_base64 = base64.b64encode(buffered.getvalue()).decode()
         
-        print(f"Success! Result size: {len(result_base64)} chars")
-        
-        return JSONResponse({
-            "success": True,
-            "saliency_map_base64": result_base64
-        })
+        return {"success": True, "saliency_map_base64": result_base64}
         
     except Exception as e:
-        print(f"Error: {str(e)}")
-        import traceback
-        traceback.print_exc()
         raise HTTPException(status_code=500, detail=str(e))
 
+# =============================================================================
+# Gradio interface (for UI)
+# =============================================================================
+
+examples = [
+    "examples/kirsten-frank-o1sXiz_LU1A-unsplash.jpg",
+    "examples/oscar-fickel-F5ze5FkEu1g-unsplash.jpg",
+    "examples/ting-tian-_79ZJS8pV70-unsplash.jpg",
+    "examples/gina-domenique-LmrAUrHinqk-unsplash.jpg",
+    "examples/robby-mccullough-r05GkQBcaPM-unsplash.jpg",
+]
+
+demo = gr.Interface(
+    fn=compute_saliency,
+    inputs=gr.Image(label="Input Image"),
+    outputs=gr.Image(label="Saliency Map"),
+    examples=examples,
+    title="Visual Saliency Prediction",
+    description="A demo to predict where humans fixate on an image using a deep learning model trained on eye movement data. Upload an image file, take a snapshot from your webcam, or paste an image from the clipboard to compute the saliency map.",
+    article="For more information on the model, check out [GitHub](https://github.com/alexanderkroner/saliency) and the corresponding [paper](https://doi.org/10.1016/j.neunet.2020.05.004).",
+    allow_flagging="never",
+    api_name="predict"
+)
+
+# Mount FastAPI to Gradio
+app = gr.mount_gradio_app(app, demo, path="/")
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)