Update app.py

app.py

@@ -1,15 +1,14 @@
 import streamlit as st
-import torch
 import cv2
-from PIL import Image
 import numpy as np
+import torch
 from torchvision import transforms
-from TranSalNet_Res import TranSalNet  # Make sure TranSalNet is accessible from your Streamlit app
+from PIL import Image
 
 # Load the model and set the device
-model = TranSalNet()
+model = TranSalNet()  # Assuming you have defined your model
 model.load_state_dict(torch.load('pretrained_models/TranSalNet_Res.pth', map_location=torch.device('cpu')))
-model.eval()  # Set the model to evaluation mode
+model.eval()
 device = torch.device('cpu')
 model.to(device)
 
@@ -27,13 +26,12 @@ if uploaded_image:
         # Preprocess the image
         img = image.resize((384, 288))
         img = np.array(img) / 255.
-        img = np.transpose(img, (2, 0, 1))
-        img = torch.from_numpy(img).unsqueeze(0).float()
-        img = img.to(device)
+        img = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)
+        img = torch.from_numpy(img)
+        img = img.type(torch.FloatTensor).to(device)
 
         # Get saliency prediction
-        with torch.no_grad():
-            pred_saliency = model(img)
+        pred_saliency = model(img)
 
         # Convert the result back to a PIL image
         toPIL = transforms.ToPILImage()
@@ -46,9 +44,50 @@ if uploaded_image:
         original_img = np.array(image)
         colorized_img = cv2.resize(colorized_img, (original_img.shape[1], original_img.shape[0]))
 
-        # You can add more post-processing here if needed
+        # Compute intensity values from the colorized image
+        intensity_map = cv2.cvtColor(colorized_img, cv2.COLOR_BGR2GRAY)
+
+        # Threshold the intensity map to create a binary mask
+        _, binary_map = cv2.threshold(intensity_map, 0, 255, cv2.THRESH_BINARY)
+
+        # Find contours in the binary map
+        contours, _ = cv2.findContours(binary_map, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        # Sort the contours by area in descending order
+        contours = sorted(contours, key=cv2.contourArea, reverse=True)
+
+        # Create an empty label map for ranking based on area
+        label_map = np.zeros_like(intensity_map)
+
+        # Rank and label each region based on area
+        for i, contour in enumerate(contours):
+            M = cv2.moments(contour)
+            if M["m00"] == 0:
+                continue
+            center_x = int(M["m10"] / M["m00"])
+            center_y = int(M["m01"] / M["m00"])
+            cv2.putText(label_map, str(i + 1), (center_x, center_y), cv2.FONT_HERSHEY_SIMPLEX, 1, 255, 2, cv2.LINE_AA)
+
+        # Blend the colorized image with the original image
+        alpha = 0.7  # Adjust the alpha value to control blending strength
+        blended_img = cv2.addWeighted(original_img, 1 - alpha, colorized_img, alpha, 0)
+
+        # Overlay the labels on the blended image
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        for i in range(1, len(contours) + 1):
+            mask = (label_map == i).astype(np.uint8)
+            x, y, w, h = cv2.boundingRect(contours[i-1])
+            org = (x, y)
+            color = (0, 0, 255)  # Red color
+            thickness = 2
+            cv2.putText(blended_img, str(i), org, font, 1, color, thickness, cv2.LINE_AA)
 
         # Display the final result
         st.image(colorized_img, caption='Colorized Saliency Map', use_column_width=True)
+        st.image(blended_img, caption='Blended Image with Labels', use_column_width=True)
+
+        # Save the final result
+        cv2.imwrite('example/result15.png', blended_img, [int(cv2.IMWRITE_JPEG_QUALITY), 200])
+        st.success('Saliency detection complete. Result saved as "example/result15.png".')
 
-st.write('Finished!')
+st.write('Finished, check the result at: example/result15.png')