Update app.py

app.py

@@ -14,6 +14,9 @@ model.load_state_dict(torch.load('pretrained_models/TranSalNet_Res.pth', map_loc
 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)
@@ -23,13 +26,7 @@ def count_and_label_red_patches(heatmap, threshold=200):
     
     original_image = np.array(image)
     
-    # Find the centroid of the red spot with the highest area
-    M_largest = cv2.moments(contours[0])
-    if M_largest["m00"] != 0:
-        cX_largest = int(M_largest["m10"] / M_largest["m00"])
-        cY_largest = int(M_largest["m01"] / M_largest["m00"])
-    else:
-        cX_largest, cY_largest = 0, 0
+    centroid_list = []  # List to store the centroids of the contours in order
     
     for i, contour in enumerate(contours, start=1):
         # Compute the centroid of the current contour
@@ -44,18 +41,24 @@ def count_and_label_red_patches(heatmap, threshold=200):
         circle_color = (0, 0, 0)  # Blue color
         cv2.circle(original_image, (cX, cY), radius, circle_color, -1)  # Draw blue circle
 
-        # Connect the current red spot to the red spot with the highest area
-        line_color = (0, 0, 0)  # Red color
-        cv2.line(original_image, (cX, cY), (cX_largest, cY_largest), line_color, 2)
-
         font = cv2.FONT_HERSHEY_SIMPLEX
         font_scale = 1
         font_color = (255, 255, 255)
         line_type = cv2.LINE_AA
         cv2.putText(original_image, str(i), (cX - 10, cY + 10), font, font_scale, font_color, 2, line_type)
+        
+        centroid_list.append((cX, cY))  # Add the centroid to the list
+
+    # Connect the red spots in the desired order
+    for i in range(len(centroid_list) - 1):
+        start_point = centroid_list[i]
+        end_point = centroid_list[i + 1]
+        line_color = (0, 0, 0)  # Red color
+        cv2.line(original_image, start_point, end_point, line_color, 2)
 
     return original_image, len(contours)
 
+
 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"])
@@ -95,6 +98,6 @@ if uploaded_image:
 
         st.image(blended_img, caption='Blended Image', use_column_width=True, channels='BGR')
 
-        # Create a dir with name example to save
+        # 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".')