Update app.py

app.py

@@ -1,22 +1,20 @@
-import streamlit as st
 import cv2
 import numpy as np
 import torch
-from torchvision import transforms, models
+from fastapi import FastAPI, UploadFile, File
 from PIL import Image
 from TranSalNet_Res import TranSalNet
-import torch.nn as nn
 from utils.data_process import preprocess_img, postprocess_img
 
+
+app = FastAPI()
+
 device = torch.device('cpu')
 model = TranSalNet()
 model.load_state_dict(torch.load('pretrained_models/TranSalNet_Res.pth', map_location=torch.device('cpu')))
 model.to(device)
 model.eval()
 
-import cv2
-import numpy as np
-
 def count_and_label_red_patches(heatmap, threshold=200):
     red_mask = heatmap[:, :, 2] > threshold
     contours, _ = cv2.findContours(red_mask.astype(np.uint8), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
@@ -58,46 +56,49 @@ def count_and_label_red_patches(heatmap, threshold=200):
 
     return original_image, len(contours)
 
+def process_image(image: Image.Image) -> np.ndarray:
+    img = image.resize((384, 288))
+    img = np.array(img)
+    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)  # Convert to BGR color space
+    img = np.array(img) / 255.
+    img = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
+    img = torch.from_numpy(img)
+    img = img.type(torch.FloatTensor).to(device)
 
-st.title('Saliency Detection App')
-st.write('Upload an image for saliency detection:')
-uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
-
-if uploaded_image:
-    image = Image.open(uploaded_image)
-    st.image(image, caption='Uploaded Image', use_column_width=True)
+    pred_saliency = model(img).squeeze().detach().numpy()
 
-    if st.button('Detect Saliency'):
-        img = image.resize((384, 288))
-        img = np.array(img)
-        img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)  # Convert to BGR color space
-        img = np.array(img) / 255.
-        img = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
-        img = torch.from_numpy(img)
-        img = img.type(torch.FloatTensor).to(device)
+    heatmap = (pred_saliency * 255).astype(np.uint8)
+    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)  # Use a blue colormap (JET)
 
-        pred_saliency = model(img).squeeze().detach().numpy()
+    heatmap = cv2.resize(heatmap, (image.width, image.height))
 
-        heatmap = (pred_saliency * 255).astype(np.uint8)
-        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)  # Use a blue colormap (JET)
+    enhanced_image = np.array(image)
+    b, g, r = cv2.split(enhanced_image)
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    b_enhanced = clahe.apply(b)
+    enhanced_image = cv2.merge((b_enhanced, g, r))
 
-        heatmap = cv2.resize(heatmap, (image.width, image.height))
+    alpha = 0.7
+    blended_img = cv2.addWeighted(enhanced_image, 1 - alpha, heatmap, alpha, 0)
 
-        enhanced_image = np.array(image)
-        b, g, r = cv2.split(enhanced_image)
-        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
-        b_enhanced = clahe.apply(b)
-        enhanced_image = cv2.merge((b_enhanced, g, r))
+    original_image, num_red_patches = count_and_label_red_patches(heatmap)
 
-        alpha = 0.7
-        blended_img = cv2.addWeighted(enhanced_image, 1 - alpha, heatmap, alpha, 0)
+    # Save processed image (optional)
+    cv2.imwrite('example/result15.png', blended_img, [int(cv2.IMWRITE_JPEG_QUALITY), 200])
 
-        original_image, num_red_patches = count_and_label_red_patches(heatmap)
+    return blended_img
 
-        st.image(original_image, caption=f'Image with {num_red_patches} Red Patches', use_column_width=True, channels='RGB')
+@app.post("/process_image")
+async def process_uploaded_image(file: UploadFile = File(...)):
+    try:
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents))
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error opening image: {str(e)}")
 
-        st.image(blended_img, caption='Blended Image', use_column_width=True, channels='BGR')
+    try:
+        processed_image = process_image(image)
+        return StreamingResponse(io.BytesIO(cv2.imencode('.png', processed_image)[1].tobytes()), media_type="image/png")
 
-        # Create a dir with the name example to save
-        cv2.imwrite('example/result15.png', blended_img, [int(cv2.IMWRITE_JPEG_QUALITY), 200])
-        st.success('Saliency detection complete. Result saved as "example/result15.png".')
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")