Update app.py

app.py

@@ -9,8 +9,6 @@ from tqdm import tqdm
 import torch.nn as nn
 from utils.data_process import preprocess_img, postprocess_img
 
-
-
 # Load the model and set the device
 device = torch.device('cpu')
 model = TranSalNet()
@@ -29,6 +27,10 @@ if uploaded_image:
 
     # Check if the user clicks a button
     if st.button('Detect Saliency'):
+        # Create a blue background image
+        blue_background = np.zeros_like(np.array(image))
+        blue_background[:] = (0, 0, 255)  # Set the background to blue (in BGR format)
+
         # Preprocess the image
         img = image.resize((384, 288))
         img = np.array(img) / 100.
@@ -48,54 +50,31 @@ if uploaded_image:
 
         # Ensure the colorized image has the same dimensions as the original image
         original_img = np.array(image)
-        colorized_img = cv2.resize(colorized_img, (original_img.shape[1], original_img.shape[0]))
-
-        # 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, 255, 0, 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)
+        colorized_img = cv2.resize(colorized_img, (original_img.shape[1], original_img.shape[0))
 
         # Create an empty label map for ranking based on area
-        label_map = np.zeros_like(intensity_map)
+        label_map = np.zeros_like(colorized_img)
 
-        # Rank and label each region based on area
+        # Overlay the labels on the blended image
+        font = cv2.FONT_HERSHEY_SIMPLEX
+        contours, _ = cv2.findContours(binary_map, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
         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)
+            cv2.putText(label_map, str(i + 1), (center_x, center_y), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 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 = (255, 0, 0)  # Blue color
-            thickness = 3
-            cv2.putText(blended_img, str(i), org, font, 1, color, thickness, cv2.LINE_AA)
-
-        
+        blended_img = cv2.addWeighted(blue_background, 1 - alpha, colorized_img, alpha, 0)
 
         # 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.success('Saliency detection complete. Result saved as "example/result15.png"')
 
 st.write('Finished, check the result at: example/result15.png')